From 90289536257ff67016b09b47a6027f19df07d6b7 Mon Sep 17 00:00:00 2001 From: William Lin Date: Sat, 8 Aug 2026 22:24:17 -0700 Subject: [PATCH] [misc] yapf pass under CI's interpreter (3.12) + pin hook language_version (#1702) --- .../scripts/clone_reference_repo.py | 8 +- .../scripts/download_hf_weights.py | 4 +- .../scripts/inspect_hf_layout.py | 44 +- .../templates/component_parity_test.py | 33 +- .../templates/family_to_diffusers.py | 57 +- .../templates/block_trace_debug.py | 82 +- .../templates/pipeline_parity_test.py | 22 +- .pre-commit-config.yaml | 1 + apps/dreamverse/dreamverse/tests/conftest.py | 1 - .../dreamverse/tests/test_config.py | 28 +- .../dreamverse/tests/test_entrypoints.py | 63 +- .../dreamverse/tests/test_gpu_pool.py | 5 +- .../dreamverse/tests/test_import_contract.py | 14 +- .../dreamverse/tests/test_mock_server.py | 309 +- .../dreamverse/tests/test_prompt_enhancer.py | 374 +- .../dreamverse/tests/test_realtime_stress.py | 342 +- .../tests/test_session_init_image.py | 24 +- .../dreamverse/tests/test_session_logging.py | 784 ++--- apps/dreamverse/scripts/modal/modal_app.py | 24 +- apps/fastvideo_studio/tests/test_gpu.py | 9 +- .../tests/test_training_config.py | 21 +- examples/inference/basic/basic.py | 25 +- .../inference/basic/basic_cosmos2_5_i2w.py | 4 +- .../inference/basic/basic_cosmos2_5_t2w.py | 6 +- .../inference/basic/basic_cosmos2_5_v2w.py | 1 - examples/inference/basic/basic_dmd.py | 19 +- examples/inference/basic/basic_dmd_new_api.py | 34 +- .../inference/basic/basic_dreamx_world.py | 7 +- examples/inference/basic/basic_flux2.py | 8 +- examples/inference/basic/basic_flux2_klein.py | 1 - examples/inference/basic/basic_flux_dev.py | 16 +- examples/inference/basic/basic_gamecraft.py | 21 +- examples/inference/basic/basic_gen3c.py | 48 +- examples/inference/basic/basic_hy15.py | 41 +- examples/inference/basic/basic_hy15_1080p.py | 27 +- examples/inference/basic/basic_hyworld.py | 2 + .../inference/basic/basic_kandinsky5_i2v.py | 6 +- .../inference/basic/basic_kandinsky5_t2v.py | 36 +- .../basic/basic_lingbotworld2_causal_fast.py | 1 - .../basic/basic_lingbotworld_base_cam.py | 8 +- examples/inference/basic/basic_longcat_i2v.py | 62 +- examples/inference/basic/basic_longcat_t2v.py | 66 +- examples/inference/basic/basic_longcat_vc.py | 80 +- examples/inference/basic/basic_ltx2.py | 27 +- .../basic/basic_ltx2_3_distilled_i2v.py | 64 +- .../basic/basic_ltx2_3_distilled_i2v_typed.py | 92 +- .../inference/basic/basic_ltx2_distilled.py | 26 +- .../basic_ltx2_distilled_fast_profile.py | 50 +- examples/inference/basic/basic_matrixgame2.py | 34 +- .../basic/basic_matrixgame2_streaming.py | 54 +- .../inference/basic/basic_minimax_h3_t2v.py | 48 +- examples/inference/basic/basic_mps.py | 32 +- examples/inference/basic/basic_ray.py | 23 +- examples/inference/basic/basic_sd35_t2i.py | 5 +- .../basic/basic_self_forcing_causal.py | 13 +- .../basic_self_forcing_causal_wan2_2_i2v.py | 12 +- .../basic_self_forcing_causal_wan2_2_t2v.py | 22 +- .../basic/basic_stable_audio_inpaint.py | 11 +- .../inference/basic/basic_turbodiffusion.py | 24 +- .../basic/basic_turbodiffusion_14b.py | 20 +- .../basic/basic_turbodiffusion_i2v.py | 5 +- examples/inference/basic/basic_wan2_2.py | 41 +- examples/inference/basic/basic_wan2_2_Fun.py | 17 +- examples/inference/basic/basic_wan2_2_i2v.py | 17 +- examples/inference/basic/basic_wan2_2_ti2v.py | 19 +- examples/inference/basic/basic_zimage.py | 4 +- examples/inference/eval/all_metrics_demo.py | 16 +- .../inference/eval/basic_ltx2_audio_eval.py | 14 +- examples/inference/eval/basic_ltx2_eval.py | 38 +- examples/inference/eval/bench_physics_iq.py | 75 +- examples/inference/eval/bench_vbench.py | 60 +- examples/inference/eval/eval_fvd.py | 27 +- examples/inference/eval/eval_ltx2_vbench.py | 84 +- .../eval/eval_third_person_separation.py | 29 +- examples/inference/eval/score_folder.py | 33 +- examples/inference/eval/score_video.py | 31 +- .../gradio/local/gradio_local_demo.py | 255 +- .../local/gradio_local_demo_ltx2_3/app.py | 90 +- .../local/gradio_local_demo_ltx2_3/config.py | 31 +- .../download_fasttext_classifiers.py | 36 +- .../gradio_local_demo_ltx2_3/examples.py | 9 +- .../prompt_enhancer.py | 73 +- .../prompt_rewrite.py | 24 +- .../gradio_local_demo_ltx2_3/rendering.py | 38 +- .../local/gradio_local_demo_ltx2_3/safety.py | 87 +- .../local/gradio_local_demo_ltx2_3/ui.py | 101 +- .../local/gradio_local_demo_matrixgame2.py | 292 +- .../gradio/serving/gradio_frontend.py | 351 +- .../gradio/serving/ray_serve_backend.py | 134 +- .../gradio/serving/start_ray_serve_app.py | 144 +- examples/inference/lora/wan_lora_inference.py | 27 +- .../lora/wan_lora_inference_from_ckpt.py | 25 +- .../inference/ltx2_3/optimized_nvfp4_t2v.py | 27 +- .../optimizations/FastWan_QAD_TAEHV.py | 103 +- .../optimizations/attention_example.py | 8 +- .../optimizations/fp4_attn_wan2_1_1_3b.py | 44 +- .../optimizations/fp8_wan2_1_1_3b.py | 75 +- .../optimizations/nvfp4_qat_wan2_1_1_3b.py | 49 +- .../inference/optimizations/qad_fp4_ab.py | 40 +- .../optimizations/spark_benchmark.py | 33 +- .../text_encoder_quant_example.py | 6 +- .../optimizations/torch_compile_example.py | 20 +- .../decode_kd_cache.py | 12 +- fastvideo-kernel/attn_qat_infer/api.py | 108 +- .../quantization/bench/bench_quant_k.py | 34 +- .../quantization/bench/bench_quant_q.py | 34 +- .../quantization/bench/bench_quant_v.py | 23 +- .../quantization/bench/bench_utils.py | 19 +- .../benchmarks/bench_fused_compress_topk.py | 71 +- fastvideo-kernel/benchmarks/bench_vsa.py | 21 +- .../benchmarks/benchmark_attn_qat_train.py | 9 +- .../python/fastvideo_kernel/__init__.py | 3 +- .../fastvideo_kernel/block_sparse_attn.py | 67 +- .../fastvideo_kernel/block_sparse_attn_256.py | 19 +- .../block_sparse_attn_cute_fwd.py | 31 +- .../block_sparse_attn_varlen.py | 71 +- .../python/fastvideo_kernel/ops.py | 91 +- .../triton_kernels/attn_qat_train.py | 865 +++-- .../block_sparse_attn_triton.py | 25 +- .../triton_kernels/fused_compress_topk.py | 34 +- .../fastvideo_kernel/triton_kernels/index.py | 103 +- .../triton_kernels/nvfp4_utils.py | 16 +- .../triton_kernels/quant_utils.py | 108 +- .../triton_kernels/sla_triton.py | 145 +- .../triton_kernels/st_attn_triton.py | 182 +- .../fastvideo_kernel/turbodiffusion_ops.py | 194 +- .../python/fastvideo_kernel/vmoba.py | 418 +-- .../python/fastvideo_kernel/vsa_utils.py | 14 +- fastvideo-kernel/tests/test_attn_qat_infer.py | 63 +- fastvideo-kernel/tests/test_attn_qat_train.py | 760 ++-- .../tests/test_fused_compress_topk.py | 88 +- fastvideo-kernel/tests/test_sta.py | 8 +- fastvideo-kernel/tests/test_turbodiffusion.py | 115 +- .../tests/test_vmoba_correctness.py | 14 +- fastvideo-kernel/tests/test_vsa.py | 167 +- fastvideo-kernel/tests/test_vsa256_forward.py | 29 +- .../tests/test_vsa256_forward_cross.py | 43 +- .../tests/test_vsa256_forward_vbs.py | 32 +- fastvideo-kernel/tests/test_vsa256_triton.py | 38 +- fastvideo-kernel/tests/test_vsa_forward.py | 61 +- fastvideo-kernel/tests/test_vsa_utils.py | 26 +- fastvideo-kernel/tests/test_vsa_varlen.py | 104 +- fastvideo-kernel/tests/utils.py | 8 +- fastvideo/dataset/__init__.py | 15 +- ...enchmark_parquet_dataset_iterable_style.py | 83 +- .../benchmark_parquet_dataset_map_style.py | 92 +- fastvideo/dataset/dataloader/parquet_io.py | 30 +- fastvideo/dataset/dataloader/record_schema.py | 88 +- fastvideo/dataset/dataloader/schema.py | 7 +- fastvideo/dataset/latent_datasets.py | 16 +- fastvideo/dataset/ltx2_precomputed_dataset.py | 47 +- .../dataset/parquet_dataset_iterable_style.py | 127 +- .../dataset/parquet_dataset_map_style.py | 145 +- fastvideo/dataset/preprocessing_datasets.py | 124 +- fastvideo/dataset/transform.py | 15 +- fastvideo/dataset/utils.py | 57 +- fastvideo/dataset/validation_dataset.py | 68 +- .../models/audio/ltx2_audio_processing.py | 6 +- fastvideo/models/audio/ltx2_audio_vae.py | 450 +-- fastvideo/models/camera/trajectory.py | 24 +- fastvideo/models/dits/base.py | 27 +- fastvideo/models/dits/causal_wanvideo.py | 468 ++- fastvideo/models/dits/cosmos.py | 269 +- fastvideo/models/dits/cosmos2_5.py | 155 +- fastvideo/models/dits/dreamx_world.py | 247 +- fastvideo/models/dits/dreamx_world_ar.py | 324 +- fastvideo/models/dits/flux.py | 63 +- fastvideo/models/dits/flux_2.py | 251 +- fastvideo/models/dits/gen3c.py | 189 +- fastvideo/models/dits/glm_image.py | 200 +- fastvideo/models/dits/hunyuangamecraft.py | 139 +- fastvideo/models/dits/hunyuanvideo.py | 270 +- fastvideo/models/dits/hunyuanvideo15.py | 280 +- fastvideo/models/dits/hyworld/__init__.py | 3 +- fastvideo/models/dits/hyworld/camera_rope.py | 10 +- fastvideo/models/dits/hyworld/data_utils.py | 2 +- fastvideo/models/dits/hyworld/hyworld.py | 196 +- fastvideo/models/dits/hyworld/pose.py | 38 +- .../models/dits/hyworld/resolution_utils.py | 27 +- .../models/dits/hyworld/retrieval_context.py | 20 +- fastvideo/models/dits/hyworld/trajectory.py | 12 +- fastvideo/models/dits/kandinsky5.py | 402 +-- .../models/dits/lingbotworld/cam_utils.py | 78 +- fastvideo/models/dits/lingbotworld/model.py | 154 +- .../models/dits/lingbotworld2/cam_utils.py | 54 +- .../models/dits/lingbotworld2/causal_fast.py | 96 +- fastvideo/models/dits/longcat.py | 244 +- fastvideo/models/dits/ltx2.py | 615 ++-- fastvideo/models/dits/magi_human.py | 163 +- .../models/dits/matrixgame2/action_module.py | 313 +- .../models/dits/matrixgame2/causal_model.py | 471 +-- fastvideo/models/dits/matrixgame2/model.py | 127 +- fastvideo/models/dits/matrixgame2/utils.py | 152 +- .../models/dits/matrixgame3/action_module.py | 183 +- fastvideo/models/dits/matrixgame3/model.py | 223 +- fastvideo/models/dits/matrixgame3/utils.py | 11 +- fastvideo/models/dits/minimax_h3.py | 38 +- fastvideo/models/dits/sd3.py | 141 +- fastvideo/models/dits/stable_audio.py | 82 +- fastvideo/models/dits/wanvideo.py | 273 +- fastvideo/models/dits/zimage.py | 55 +- fastvideo/models/encoders/base.py | 26 +- fastvideo/models/encoders/bert.py | 9 +- fastvideo/models/encoders/clip.py | 182 +- fastvideo/models/encoders/gemma.py | 211 +- .../models/encoders/glm_image_ar_loader.py | 19 +- fastvideo/models/encoders/lingbotworld2_t5.py | 38 +- fastvideo/models/encoders/llama.py | 92 +- .../models/encoders/minimax_h3_qwen3_vl.py | 76 +- fastvideo/models/encoders/mistral3.py | 6 +- fastvideo/models/encoders/qwen2_5.py | 141 +- .../models/encoders/qwen2_5_vl_custom.py | 287 +- fastvideo/models/encoders/qwen3.py | 93 +- fastvideo/models/encoders/reason1.py | 161 +- fastvideo/models/encoders/siglip.py | 76 +- .../encoders/stable_audio_conditioner.py | 39 +- fastvideo/models/encoders/t5.py | 217 +- fastvideo/models/encoders/t5_hf.py | 16 +- fastvideo/models/encoders/t5gemma.py | 12 +- fastvideo/models/encoders/vision.py | 5 +- fastvideo/models/hf_transformer_utils.py | 21 +- .../benchmarks/benchmark_weight_loading.py | 65 +- .../benchmark_weight_loading_comparison.py | 28 +- fastvideo/models/loader/component_loader.py | 352 +- fastvideo/models/loader/utils.py | 15 +- fastvideo/models/loader/weight_utils.py | 111 +- fastvideo/models/mask_utils.py | 18 +- fastvideo/models/parameter.py | 80 +- fastvideo/models/registry.py | 86 +- fastvideo/models/schedulers/base.py | 8 +- .../scheduling_flow_map_euler_discrete.py | 26 +- .../scheduling_flow_match_euler_discrete.py | 138 +- .../scheduling_flow_unipc_multistep.py | 157 +- .../schedulers/scheduling_minimax_h3.py | 15 +- fastvideo/models/schedulers/scheduling_rcm.py | 35 +- .../scheduling_self_forcing_flow_match.py | 64 +- .../schedulers/scheduling_unipc_multistep.py | 75 +- fastvideo/models/upsamplers/hunyuan15.py | 6 +- fastvideo/models/upsamplers/ltx2_upsampler.py | 3 +- fastvideo/models/utils.py | 33 +- fastvideo/models/vaes/autoencoder_kl.py | 2 +- fastvideo/models/vaes/common.py | 152 +- fastvideo/models/vaes/cosmos25wanvae.py | 46 +- fastvideo/models/vaes/flux2_components.py | 181 +- fastvideo/models/vaes/flux2vae.py | 115 +- fastvideo/models/vaes/gamecraftvae.py | 102 +- fastvideo/models/vaes/gamecraftvae_blocks.py | 83 +- fastvideo/models/vaes/gen3c_tokenizer_vae.py | 32 +- fastvideo/models/vaes/hunyuan15vae.py | 98 +- fastvideo/models/vaes/hunyuanvae.py | 230 +- fastvideo/models/vaes/hyworldvae.py | 170 +- fastvideo/models/vaes/lingbotworld2_wanvae.py | 206 +- fastvideo/models/vaes/ltx2vae.py | 205 +- fastvideo/models/vaes/minimax_h3_audio.py | 61 +- fastvideo/models/vaes/minimax_h3_video.py | 133 +- fastvideo/models/vaes/oobleck.py | 105 +- fastvideo/models/vaes/sa_audio.py | 6 +- fastvideo/models/vaes/wanvae.py | 225 +- fastvideo/models/vision_utils.py | 70 +- .../api/test_attention_selector_resolution.py | 4 +- .../test_attn_qat_infer_capability_gate.py | 6 +- fastvideo/tests/api/test_cli_translation.py | 48 +- .../tests/api/test_extra_overrides_routing.py | 8 +- fastvideo/tests/api/test_ltx2_continuation.py | 50 +- .../api/test_ltx2_gpu_pool_translation.py | 31 +- .../tests/api/test_ltx2_param_mapping.py | 56 +- .../tests/api/test_ltx2_stage_overrides.py | 9 +- fastvideo/tests/api/test_overrides.py | 35 +- .../tests/api/test_schema_parity_inventory.py | 25 +- fastvideo/tests/api/test_typed_quant_flow.py | 20 +- fastvideo/tests/api/test_validation_errors.py | 38 +- .../test_attn_qat_infer_arch_gate.py | 8 +- .../test_flash_attn_cute_custom_op.py | 128 +- .../test_flash_attn_default_custom_op.py | 24 +- .../test_flash_attn_no_pad_custom_op.py | 19 +- .../test_flash_attn_no_pad_resolver.py | 5 +- .../attention/test_selector_role_override.py | 6 +- fastvideo/tests/audio/test_ltx2_bwe.py | 60 +- fastvideo/tests/conftest.py | 3 +- .../tests/contract/test_ci_test_collection.py | 30 +- .../tests/contract/test_dreamverse_shape.py | 38 +- fastvideo/tests/contract/test_dynamo_shape.py | 49 +- .../tests/contract/test_generate_async.py | 26 +- .../contract/test_ltx2_compile_conditions.py | 5 +- .../tests/contract/test_modal_fa4_policy.py | 11 +- .../tests/dataset/test_load_image_retry.py | 6 +- .../dataset/test_parquet_dataset_map_style.py | 8 +- fastvideo/tests/dataset/test_parquet_io.py | 2 - .../tests/distributed/test_sp_hunyuanvideo.py | 41 +- fastvideo/tests/distributed/test_sp_ltx2.py | 48 +- fastvideo/tests/distributed/test_sp_wan.py | 44 +- fastvideo/tests/encoders/test_clip_encoder.py | 11 +- .../test_gemma_connector_register_padding.py | 16 +- fastvideo/tests/encoders/test_hyt5_encoder.py | 31 +- .../tests/encoders/test_llama_encoder.py | 17 +- .../tests/encoders/test_qwen2_5_encoder.py | 75 +- .../tests/encoders/test_siglip_encoder.py | 67 +- fastvideo/tests/encoders/test_t5_encoder.py | 53 +- .../entrypoints/streaming/test_auxiliaries.py | 41 +- .../entrypoints/streaming/test_gpu_pool.py | 34 +- .../streaming/test_health_routes.py | 34 +- .../streaming/test_prompt_enhancer.py | 9 +- .../streaming/test_prompt_providers.py | 64 +- .../entrypoints/streaming/test_protocol.py | 18 +- .../entrypoints/streaming/test_router.py | 64 +- .../entrypoints/streaming/test_server.py | 40 +- .../entrypoints/streaming/test_session.py | 18 +- .../streaming/test_session_init_image.py | 38 +- .../streaming/test_session_store.py | 12 +- .../streaming/test_stream_encoder.py | 23 +- .../entrypoints/streaming/test_worker.py | 5 +- .../tests/entrypoints/test_openai_api.py | 78 +- .../test_openai_api_integration.py | 25 +- .../tests/entrypoints/test_video_generator.py | 8 +- fastvideo/tests/eval/conftest.py | 21 +- .../eval/test_metric_identity_invariants.py | 12 +- .../eval/test_metric_score_regression.py | 48 +- fastvideo/tests/golden_gate/_harness.py | 26 +- fastvideo/tests/golden_gate/test_dreamx.py | 46 +- fastvideo/tests/golden_gate/test_flux.py | 4 +- .../tests/golden_gate/test_flux2_klein.py | 12 +- fastvideo/tests/golden_gate/test_gamecraft.py | 15 +- fastvideo/tests/golden_gate/test_gen3c.py | 12 +- fastvideo/tests/golden_gate/test_glm_image.py | 9 +- .../tests/golden_gate/test_kandinsky5.py | 3 +- fastvideo/tests/golden_gate/test_lingbot.py | 8 +- fastvideo/tests/golden_gate/test_longcat.py | 2 +- fastvideo/tests/golden_gate/test_ltx2.py | 98 +- .../tests/golden_gate/test_matrixgame.py | 55 +- .../tests/golden_gate/test_minimax_h3_t2v.py | 20 +- fastvideo/tests/golden_gate/test_sd35.py | 3 +- .../tests/golden_gate/test_stable_audio.py | 2 +- fastvideo/tests/golden_gate/test_wan_t2v.py | 8 +- fastvideo/tests/golden_gate/test_zimage.py | 3 +- fastvideo/tests/hooks/test_hooks_api.py | 13 +- .../tests/hooks/test_layerwise_offload.py | 122 +- .../tests/inference/bsa/test_bsa_inference.py | 8 +- .../lora/test_lora_inference_similarity.py | 123 +- fastvideo/tests/inference/test_cfg_gating.py | 7 +- .../inference/vmoba/test_vmoba_inference.py | 3 +- .../layers/test_rmsnorm_forward_dispatch.py | 29 +- .../lora_extraction/test_lora_extraction.py | 16 +- fastvideo/tests/modal/kernel_build_cache.py | 24 +- fastvideo/tests/modal/kernel_cache_smoke.py | 34 +- fastvideo/tests/modal/launch_l40s_job.py | 47 +- fastvideo/tests/modal/pr_test.py | 256 +- fastvideo/tests/modal/ssim_test.py | 139 +- .../tests/modal/test_kernel_build_cache.py | 38 +- fastvideo/tests/modal/test_pr_test.py | 59 +- fastvideo/tests/modal/test_ssim_test.py | 2 + .../nightly/test_e2e_dmd_t2v_crush_smol.py | 192 +- .../test_e2e_i2v_overfit_single_sample.py | 171 +- .../test_e2e_kandinsky5_dmd_t2v_overfit.py | 105 +- .../test_e2e_ltx2_overfit_new_stack.py | 59 +- .../nightly/test_e2e_overfit_single_sample.py | 171 +- .../tests/ops/quantization/test_absmax_fp8.py | 10 +- .../ops/quantization/test_nvfp4_config.py | 12 +- .../quantization/test_nvfp4_ltx2_wiring.py | 33 +- .../ops/quantization/test_nvfp4_purge.py | 3 +- .../tests/ops/test_rotary_pos_embed_cache.py | 41 +- .../tests/performance/compare_baseline.py | 86 +- fastvideo/tests/performance/dashboard.py | 38 +- fastvideo/tests/performance/identity.py | 48 +- fastvideo/tests/performance/seed_baseline.py | 101 +- .../test_benchmark_config_validation.py | 10 +- .../test_compare_baseline_policy.py | 138 +- .../tests/performance/test_dashboard_api.py | 174 +- .../performance/test_dashboard_plotly.py | 28 +- .../performance/test_dashboard_service.py | 43 +- fastvideo/tests/performance/test_identity.py | 55 +- .../performance/test_inference_performance.py | 128 +- ...t_inference_performance_component_times.py | 24 +- .../test_inference_performance_identity.py | 8 +- ...est_inference_performance_result_schema.py | 5 +- fastvideo/tests/ssim/bootstrap_references.py | 15 +- fastvideo/tests/ssim/conftest.py | 17 +- .../tests/ssim/latent_similarity_utils.py | 105 +- fastvideo/tests/ssim/reference_utils.py | 4 +- fastvideo/tests/ssim/reference_videos_cli.py | 152 +- .../tests/ssim/test_causal_similarity.py | 3 +- .../ssim/test_dreamx_world_similarity.py | 3 +- fastvideo/tests/ssim/test_flux2_similarity.py | 10 +- .../tests/ssim/test_flux_t2i_similarity.py | 9 +- .../tests/ssim/test_gamecraft_similarity.py | 73 +- fastvideo/tests/ssim/test_gen3c_similarity.py | 49 +- .../tests/ssim/test_lingbot_similarity.py | 115 +- .../tests/ssim/test_longcat_similarity.py | 125 +- fastvideo/tests/ssim/test_ltx2_similarity.py | 12 +- .../tests/ssim/test_matrixgame2_similarity.py | 10 +- .../tests/ssim/test_matrixgame3_similarity.py | 5 +- .../tests/ssim/test_reference_videos_cli.py | 20 +- fastvideo/tests/ssim/test_sd35_similarity.py | 6 +- .../ssim/test_stable_audio_similarity.py | 3 +- .../ssim/test_turbodiffusion_similarity.py | 6 +- .../tests/ssim/test_wan_i2v_similarity.py | 61 +- .../tests/ssim/test_wan_t2v_similarity.py | 61 +- ...st_kandinsky5_dmd_stage_backend_engages.py | 7 +- .../stages/test_kandinsky5_target_dtype.py | 10 +- .../test_magi_human_sr_latent_preparation.py | 6 +- fastvideo/tests/stages/test_text_encoding.py | 64 +- fastvideo/tests/test_partial_hf_download.py | 18 +- .../tests/train/callbacks/test_callback.py | 84 +- fastvideo/tests/train/callbacks/test_ema.py | 37 +- .../tests/train/callbacks/test_grad_clip.py | 26 +- .../tests/train/callbacks/test_validation.py | 82 +- .../train/methods/grad_norm_regression.py | 38 +- .../train/methods/test_cosmos_finetune.py | 34 +- .../tests/train/methods/test_ltx2_finetune.py | 61 +- .../methods/test_matrixgame2_finetune.py | 52 +- .../train/methods/test_minimax_h3_finetune.py | 84 +- .../methods/test_streaming_long_tuning.py | 59 +- .../tests/train/methods/test_wan_causal_cd.py | 51 +- .../train/methods/test_wan_causal_dfsft.py | 38 +- .../test_wan_causal_dfsft_framewise.py | 34 +- .../train/methods/test_wan_causal_tfsft.py | 40 +- .../tests/train/methods/test_wan_finetune.py | 34 +- .../test_kandinsky5_qat_attention_engages.py | 20 +- .../tests/train/models/test_load_cosmos.py | 14 +- .../tests/train/models/test_load_hunyuan.py | 9 +- .../train/models/test_load_kandinsky5.py | 10 +- .../tests/train/models/test_load_longcat.py | 14 +- .../tests/train/models/test_load_ltx2.py | 43 +- .../train/models/test_load_matrixgame2.py | 14 +- fastvideo/tests/train/models/test_load_wan.py | 14 +- .../train/models/test_load_wan_causal.py | 9 +- .../tests/train/trainer/test_validation.py | 8 +- .../tests/train/utils/test_checkpoint.py | 13 +- fastvideo/tests/train/utils/test_config.py | 30 +- .../test_finetune_config_construction.py | 11 +- .../train/utils/test_validation_media.py | 5 +- .../training/VSA/test_training_loss_VSA.py | 97 +- .../tests/training/Vanilla/mfu_calculation.py | 165 +- .../training/Vanilla/test_training_loss.py | 99 +- .../checkpoint/test_checkpointing_utils.py | 3 +- .../training/distill/test_anyflow_onpolicy.py | 27 +- .../training/distill/test_anyflow_pretrain.py | 67 +- .../training/distill/test_anyflow_smoke.py | 85 +- .../training/distill/test_distill_dmd.py | 92 +- .../tests/training/lora/test_lora_training.py | 91 +- .../self-forcing/test_self_forcing.py | 199 +- fastvideo/tests/training/test_trackers.py | 2 + fastvideo/tests/transformers/test_cosmos.py | 42 +- .../tests/transformers/test_cosmos2_5.py | 276 +- fastvideo/tests/transformers/test_flux.py | 34 +- .../transformers/test_hunyuangamecraft.py | 28 +- fastvideo/tests/transformers/test_hyworld.py | 49 +- fastvideo/tests/transformers/test_wanvideo.py | 25 +- fastvideo/tests/utils.py | 9 +- fastvideo/tests/vaes/test_hunyuan15_vae.py | 31 +- fastvideo/tests/vaes/test_wan_vae.py | 18 +- fastvideo/tests/vaes/test_wan_vae_compile.py | 3 +- .../test_kandinsky5_overfit_caption.py | 4 +- .../test_minimax_h3_overfit_preprocess.py | 28 +- .../eval/glmasr/modeling_glmasr.py | 45 +- .../eval/glmasr/processing_glmasr.py | 44 +- fastvideo/third_party/eval/synchformer/ast.py | 22 +- .../eval/synchformer/hf_src/modeling_ast.py | 72 +- .../eval/synchformer/motionformer.py | 42 +- .../third_party/eval/synchformer/resnet.py | 53 +- .../eval/synchformer/synchformer.py | 9 +- .../eval/synchformer/transformer.py | 12 +- .../third_party/eval/synchformer/utils.py | 69 +- .../eval/synchformer/video_model_builder.py | 21 +- .../eval/synchformer/vit_helper.py | 67 +- .../block_sparse_attention/bsa_interface.py | 484 +-- .../block_sparse_attention/common.py | 59 +- .../block_sparse_attention/communicate.py | 22 +- .../flash_attn_bsa_varlen_mask.py | 658 ++-- fastvideo/third_party/pynvml.py | 3054 ++++++++++------- scripts/check_docs_links.py | 8 +- .../convert_flux2_klein.py | 36 +- .../convert_gamecraft_full.py | 149 +- .../convert_gamecraft_vae.py | 4 +- .../convert_gamecraft_weights.py | 99 +- .../convert_gen3c_to_fastvideo.py | 133 +- .../convert_lingbotworld2_causal_fast.py | 7 +- .../convert_ltx2_weights.py | 37 +- .../convert_magi_human_to_diffusers.py | 50 +- ...convert_turbodiffusion_i2v_to_diffusers.py | 58 +- .../convert_turbodiffusion_to_diffusers.py | 65 +- .../checkpoint_conversion/create_hf_repo.py | 93 +- .../dreamx_world_to_diffusers.py | 11 +- .../extract_llava_text_encoder.py | 94 +- .../longcat_to_fastvideo.py | 240 +- .../pt_to_safetensors.py | 61 +- .../push_magi_human_to_hf.py | 16 +- .../stable_audio_to_diffusers.py | 19 +- .../validate_longcat_weights.py | 96 +- .../checkpoint_conversion/wan_to_diffusers.py | 189 +- .../generate_validation_json.py | 8 +- .../dataset_preparation/prepare_json_file.py | 38 +- scripts/demo_anyflow_14b.py | 46 +- scripts/lora_extraction/extract_lora.py | 27 +- .../lora_inference_comparison.py | 30 +- scripts/lora_extraction/merge_lora.py | 114 +- scripts/lora_extraction/verify_lora.py | 57 +- scripts/ltx2_sr_alignment.py | 58 +- scripts/verify_anyflow_fastvideo_parity.py | 56 +- ...test_dreamx_world_ar_transformer_parity.py | 57 +- ...dreamx_world_camera_conditioning_parity.py | 10 +- .../test_dreamx_world_pipeline_config.py | 40 +- .../test_dreamx_world_scheduler_parity.py | 3 +- .../test_dreamx_world_text_encoder_parity.py | 11 +- .../test_dreamx_world_transformer_parity.py | 75 +- .../test_dreamx_world_vae_parity.py | 1 + .../test_minimax_h3_qwen3_vl_parity.py | 27 +- .../flux/test_flux_dev_component_loaders.py | 43 +- .../flux2/debug_flux2_klein_tp_trace.py | 91 +- .../generate_flux2_full_image_compare.py | 32 +- .../flux2/test_flux2_component_parity.py | 169 +- .../test_gamecraft_encoders_parity.py | 119 +- .../gamecraft/test_gamecraft_parity.py | 311 +- .../test_gamecraft_pipeline_parity.py | 239 +- .../gamecraft/test_gamecraft_vae_parity.py | 52 +- tests/local_tests/gen3c/test_gen3c.py | 224 +- .../gen3c/test_gen3c_pipeline_smoke.py | 20 +- tests/local_tests/glm_image/test_ar_parity.py | 37 +- .../glm_image/test_edit_pipeline_parity.py | 15 +- .../glm_image/test_pipeline_parity.py | 124 +- tests/local_tests/glm_image/test_t5_parity.py | 58 +- .../glm_image/test_transformer_parity.py | 73 +- .../glm_image/test_transformer_strict_load.py | 10 +- .../local_tests/glm_image/test_vae_parity.py | 13 +- .../helpers/magi_human_upstream.py | 33 +- ...test_kandinsky5_lite_transformer_parity.py | 33 +- ...est_lingbotworld2_causal_fast_preflight.py | 6 +- tests/local_tests/ltx2/test_ltx2.py | 113 +- tests/local_tests/ltx2/test_ltx2_audio.py | 87 +- tests/local_tests/ltx2/test_ltx2_audio_vae.py | 146 +- .../ltx2/test_ltx2_gemma_encoder.py | 29 +- .../ltx2/test_ltx2_gemma_parity.py | 66 +- .../ltx2/test_ltx2_pipeline_smoke.py | 31 +- tests/local_tests/ltx2/test_ltx2_registry.py | 17 +- tests/local_tests/ltx2/test_ltx2_vae.py | 64 +- .../ltx2/test_ltx2_vae_official.py | 30 +- .../test_magi_human_distill_parity.py | 86 +- .../magi_human/test_magi_human_parity.py | 79 +- .../test_magi_human_pipeline_parity.py | 256 +- .../test_magi_human_pipeline_smoke.py | 22 +- ...est_magi_human_sa_audio_official_parity.py | 27 +- .../test_magi_human_sa_audio_parity.py | 28 +- ...test_magi_human_sr1080p_pipeline_parity.py | 25 +- .../test_magi_human_sr540p_pipeline_parity.py | 29 +- .../test_magi_human_t5gemma_parity.py | 53 +- .../test_magi_human_ti2v_pipeline_parity.py | 17 +- .../magi_human/test_magi_human_vae_parity.py | 39 +- tests/local_tests/minimax_h3/_reference.py | 16 +- .../minimax_h3/test_minimax_h3_packing.py | 7 +- .../test_minimax_h3_ref2va_media.py | 10 +- .../test_minimax_h3_ref2va_packing.py | 3 +- .../test_dreamx_world_pipeline_parity.py | 11 +- .../test_dreamx_world_pipeline_smoke.py | 8 +- .../pipelines/test_flux2_pipeline_parity.py | 193 +- .../pipelines/test_flux2_pipeline_smoke.py | 32 +- .../test_flux_dev_pipeline_parity.py | 9 +- .../pipelines/test_flux_dev_pipeline_smoke.py | 9 +- ...ngbotworld2_causal_fast_pipeline_parity.py | 7 +- ...ingbotworld2_causal_fast_pipeline_smoke.py | 1 - .../test_minimax_h3_pipeline_smoke.py | 5 +- .../pipelines/test_zimage_pipeline_parity.py | 14 +- .../pipelines/test_zimage_pipeline_smoke.py | 1 - .../sd35/test_sd35_component_parity.py | 34 +- .../test_oobleck_vae_official_parity.py | 72 +- .../stable_audio/test_oobleck_vae_parity.py | 72 +- .../test_stable_audio_a2a_parity.py | 15 +- .../test_stable_audio_inpaint_parity.py | 11 +- .../test_stable_audio_pipeline_parity.py | 43 +- .../test_stable_audio_pipeline_smoke.py | 7 +- tests/local_tests/test_nvfp4_fa4.py | 44 +- tests/local_tests/test_train_rewards.py | 10 +- tests/local_tests/test_train_rl_sampling.py | 2 +- .../test_trainer_managed_optimization.py | 8 +- .../test_minimax_h3_transformer_parity.py | 8 +- .../vaes/test_minimax_h3_audio_vae_parity.py | 4 +- .../vaes/test_minimax_h3_video_vae_parity.py | 14 +- .../wan22/test_i2v_record_no_clip.py | 17 +- .../zimage/test_zimage_encoder_parity.py | 145 +- .../zimage/test_zimage_scheduler_parity.py | 23 +- .../zimage/test_zimage_tokenizer_parity.py | 11 +- .../zimage/test_zimage_transformer_parity.py | 17 +- .../zimage/test_zimage_vae_parity.py | 9 +- tests/test_bsa.py | 173 +- 582 files changed, 19387 insertions(+), 23526 deletions(-) diff --git a/.agents/skills/add-model-01-prep/scripts/clone_reference_repo.py b/.agents/skills/add-model-01-prep/scripts/clone_reference_repo.py index b2c898a8b..3bc2e27e1 100755 --- a/.agents/skills/add-model-01-prep/scripts/clone_reference_repo.py +++ b/.agents/skills/add-model-01-prep/scripts/clone_reference_repo.py @@ -10,9 +10,7 @@ from pathlib import Path def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Clone a reference repo for FastVideo parity tests." - ) + parser = argparse.ArgumentParser(description="Clone a reference repo for FastVideo parity tests.") parser.add_argument("repo_url", help="Official reference repository URL") parser.add_argument("target_dir", help="Directory to clone into") parser.add_argument("--branch", help="Branch or tag to clone") @@ -62,9 +60,7 @@ def gitignore_entry_for(target: Path) -> str: try: relative = resolved.relative_to(root) except ValueError as exc: - raise ValueError( - "--update-gitignore requires target_dir to be under the current directory" - ) from exc + raise ValueError("--update-gitignore requires target_dir to be under the current directory") from exc text = relative.as_posix().rstrip("/") return "/" + text + "/" diff --git a/.agents/skills/add-model-01-prep/scripts/download_hf_weights.py b/.agents/skills/add-model-01-prep/scripts/download_hf_weights.py index a5e2dbd8f..1890aaad6 100755 --- a/.agents/skills/add-model-01-prep/scripts/download_hf_weights.py +++ b/.agents/skills/add-model-01-prep/scripts/download_hf_weights.py @@ -8,14 +8,12 @@ import os import sys from pathlib import Path - HF_TOKEN_ENV_KEYS = ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( - description="Download a HF model snapshot or selected files into a local directory." - ) + description="Download a HF model snapshot or selected files into a local directory.") parser.add_argument("repo_id", help="HF repo id, for example Org/Model") parser.add_argument("local_dir", help="Destination directory") parser.add_argument("--repo-type", default="model", help="HF repo type (default: model)") diff --git a/.agents/skills/add-model-01-prep/scripts/inspect_hf_layout.py b/.agents/skills/add-model-01-prep/scripts/inspect_hf_layout.py index 3017b5ee0..e7280ec6b 100755 --- a/.agents/skills/add-model-01-prep/scripts/inspect_hf_layout.py +++ b/.agents/skills/add-model-01-prep/scripts/inspect_hf_layout.py @@ -10,7 +10,6 @@ import sys from pathlib import Path from typing import Any - HF_TOKEN_ENV_KEYS = ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY") RAW_WEIGHT_SUFFIXES = (".safetensors", ".pt", ".pth", ".ckpt", ".bin") KNOWN_COMPONENTS = { @@ -34,8 +33,7 @@ KNOWN_COMPONENTS = { def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( - description="Classify a HF repo or local directory as Diffusers, raw, custom, or unknown." - ) + description="Classify a HF repo or local directory as Diffusers, raw, custom, or unknown.") parser.add_argument("source", help="HF repo id or local weights directory") parser.add_argument("--repo-type", default="model", help="HF repo type (default: model)") parser.add_argument("--revision", help="HF revision to inspect") @@ -94,14 +92,12 @@ def load_remote_files( ) -> list[str]: from huggingface_hub import list_repo_files - return sorted( - list_repo_files( - repo_id, - repo_type=repo_type, - revision=revision, - token=token, - ) - ) + return sorted(list_repo_files( + repo_id, + repo_type=repo_type, + revision=revision, + token=token, + )) def load_remote_model_index( @@ -215,24 +211,24 @@ def build_result(args: argparse.Namespace) -> dict[str, Any]: "components_seen": components, "file_count": len(files), "file_scan_truncated": truncated, - "files_sample": files[: args.sample_limit], + "files_sample": files[:args.sample_limit], } def print_human(result: dict[str, Any]) -> None: for key in ( - "source", - "source_kind", - "repo_type", - "revision", - "token_env", - "source_layout", - "needs_conversion", - "model_index_class", - "model_index_diffusers_version", - "model_index_error", - "file_count", - "file_scan_truncated", + "source", + "source_kind", + "repo_type", + "revision", + "token_env", + "source_layout", + "needs_conversion", + "model_index_class", + "model_index_diffusers_version", + "model_index_error", + "file_count", + "file_scan_truncated", ): value = result.get(key) if value is not None: diff --git a/.agents/skills/add-model-02-parity/templates/component_parity_test.py b/.agents/skills/add-model-02-parity/templates/component_parity_test.py index fbe7afbe1..8d24e292a 100644 --- a/.agents/skills/add-model-02-parity/templates/component_parity_test.py +++ b/.agents/skills/add-model-02-parity/templates/component_parity_test.py @@ -18,7 +18,6 @@ import pytest import torch from torch.testing import assert_close - os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29519") os.environ.setdefault("DISABLE_SP", "1") @@ -35,15 +34,10 @@ FASTVIDEO_CONFIG_CLASS = "" # TODO. FASTVIDEO_MODEL_MODULE = "fastvideo.models.." # TODO. FASTVIDEO_MODEL_CLASS = "" # TODO. -OFFICIAL_REF_DIR = Path( - os.getenv("_OFFICIAL_REF_DIR", REPO_ROOT / "") -) -LOCAL_WEIGHTS_DIR = Path( - os.getenv("_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / FAMILY) -) -CONVERTED_WEIGHTS_DIR = Path( - os.getenv("_CONVERTED_WEIGHTS_DIR", REPO_ROOT / "converted_weights" / FAMILY) -) +OFFICIAL_REF_DIR = Path(os.getenv("_OFFICIAL_REF_DIR", REPO_ROOT / "")) +LOCAL_WEIGHTS_DIR = Path(os.getenv("_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / FAMILY)) +CONVERTED_WEIGHTS_DIR = Path(os.getenv("_CONVERTED_WEIGHTS_DIR", + REPO_ROOT / "converted_weights" / FAMILY)) def _resolve_hf_token() -> str | None: @@ -99,18 +93,14 @@ def _load_official_model(device: torch.device, dtype: torch.dtype) -> torch.nn.M model = OfficialClass() # TODO: pass official config kwargs. state_dict = {} # TODO: load official state dict from LOCAL_WEIGHTS_DIR. missing, unexpected = model.load_state_dict(state_dict, strict=True) - assert not missing and not unexpected, ( - f"official load mismatch missing={missing[:5]} unexpected={unexpected[:5]}" - ) + assert not missing and not unexpected, (f"official load mismatch missing={missing[:5]} unexpected={unexpected[:5]}") return model.to(device=device, dtype=dtype).eval() def _load_fastvideo_model(device: torch.device, dtype: torch.dtype) -> torch.nn.Module: """Load the FastVideo component with the same tensor content.""" if not CONVERTED_WEIGHTS_DIR.exists() and not LOCAL_WEIGHTS_DIR.exists(): - pytest.skip( - f"No FastVideo loadable weights: {CONVERTED_WEIGHTS_DIR} or {LOCAL_WEIGHTS_DIR}" - ) + pytest.skip(f"No FastVideo loadable weights: {CONVERTED_WEIGHTS_DIR} or {LOCAL_WEIGHTS_DIR}") # TODO: replace with the bucket-specific FastVideo config/class/loader. # DiT examples: @@ -127,8 +117,7 @@ def _load_fastvideo_model(device: torch.device, dtype: torch.dtype) -> torch.nn. state_dict = {} # TODO: load converted or directly mapped state dict. missing, unexpected = model.load_state_dict(state_dict, strict=True) assert not missing and not unexpected, ( - f"FastVideo load mismatch missing={missing[:5]} unexpected={unexpected[:5]}" - ) + f"FastVideo load mismatch missing={missing[:5]} unexpected={unexpected[:5]}") return model.to(device=device, dtype=dtype).eval() @@ -187,11 +176,9 @@ def test_component_parity(): assert official_out.shape == fastvideo_out.shape diff = (official_out - fastvideo_out).abs() - print( - f"official abs_mean={official_out.abs().mean().item():.6f} " - f"fastvideo abs_mean={fastvideo_out.abs().mean().item():.6f} " - f"diff_max={diff.max().item():.6f} diff_mean={diff.mean().item():.6f}" - ) + print(f"official abs_mean={official_out.abs().mean().item():.6f} " + f"fastvideo abs_mean={fastvideo_out.abs().mean().item():.6f} " + f"diff_max={diff.max().item():.6f} diff_mean={diff.mean().item():.6f}") # TODO: pick tolerance by scope: # - single block / same kernel: 1e-4 diff --git a/.agents/skills/add-model-07-conversion/templates/family_to_diffusers.py b/.agents/skills/add-model-07-conversion/templates/family_to_diffusers.py index fa04b1158..0d12589bb 100644 --- a/.agents/skills/add-model-07-conversion/templates/family_to_diffusers.py +++ b/.agents/skills/add-model-07-conversion/templates/family_to_diffusers.py @@ -27,7 +27,6 @@ try: except ImportError: # pragma: no cover - optional local conversion dependency snapshot_download = None - # TODO: fill with authoritative component prefixes for monolithic checkpoints. # Example: {"model.model.": "transformer", "pretransform.model.": "vae"} COMPONENT_PREFIXES: dict[str, str] = {} @@ -47,10 +46,7 @@ SKIP_PATTERNS: tuple[str, ...] = () def _hf_token() -> str | None: - return ( - os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") - or os.environ.get("HF_API_KEY") - ) + return (os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get("HF_API_KEY")) def resolve_src(src: str, revision: str | None) -> Path: @@ -95,11 +91,10 @@ def apply_mapping(key: str) -> str | None: return key -def split_monolithic( - state: dict[str, torch.Tensor], -) -> dict[str, OrderedDict[str, torch.Tensor]]: +def split_monolithic(state: dict[str, torch.Tensor], ) -> dict[str, OrderedDict[str, torch.Tensor]]: components: dict[str, OrderedDict[str, torch.Tensor]] = { - name: OrderedDict() for name in set(COMPONENT_PREFIXES.values()) + name: OrderedDict() + for name in set(COMPONENT_PREFIXES.values()) } intentionally_skipped: list[str] = [] unowned: list[str] = [] @@ -117,10 +112,8 @@ def split_monolithic( unowned.append(key) if unowned: sample = ", ".join(unowned[:10]) - raise ValueError( - f"Unowned monolithic keys: {len(unowned)}. " - f"Add COMPONENT_PREFIXES or SKIP_PATTERNS entries. Sample: {sample}" - ) + raise ValueError(f"Unowned monolithic keys: {len(unowned)}. " + f"Add COMPONENT_PREFIXES or SKIP_PATTERNS entries. Sample: {sample}") if intentionally_skipped: print(f"Intentionally skipped {len(intentionally_skipped)} keys") return {name: weights for name, weights in components.items() if weights} @@ -143,8 +136,12 @@ def build_component_configs(_src_dir: Path) -> dict[str, dict[str, Any]]: # TODO: emit config content accepted by FastVideo loaders. Most components use # config.json; schedulers use scheduler_config.json. return { - "transformer": {"_class_name": ""}, - "vae": {"_class_name": ""}, + "transformer": { + "_class_name": "" + }, + "vae": { + "_class_name": "" + }, } @@ -177,19 +174,13 @@ def build_model_index( } if revision: index["_fastvideo_converted_revision"] = revision - return { - key: value - for key, value in index.items() - if key.startswith("_") or key in available_components - } + return {key: value for key, value in index.items() if key.startswith("_") or key in available_components} def validate_component_configs(configs: dict[str, dict[str, Any]]) -> None: # TODO: instantiate each FastVideo config and call update_model_arch(...) or # update_model_config(...) with this JSON so unknown emitted keys fail here. - placeholder_configs = [ - name for name, config in configs.items() if "<" in json.dumps(config) - ] + placeholder_configs = [name for name, config in configs.items() if "<" in json.dumps(config)] if placeholder_configs: raise ValueError(f"Replace config placeholders for: {placeholder_configs}") @@ -201,9 +192,7 @@ def verify_conversion( del dst_dir, components # TODO: load each emitted stateful component through its production loader and # assert strict load, or document exact allowed missing/unexpected keys. - raise NotImplementedError( - "Implement production config validation and strict-load checks" - ) + raise NotImplementedError("Implement production config validation and strict-load checks") def write_component( @@ -216,9 +205,7 @@ def write_component( if component_dir.exists() and any(component_dir.iterdir()): shutil.rmtree(component_dir) component_dir.mkdir(parents=True, exist_ok=True) - save_file( - dict(state), str(component_dir / "diffusion_pytorch_model.safetensors") - ) + save_file(dict(state), str(component_dir / "diffusion_pytorch_model.safetensors")) if config is not None: config_path = component_dir / config_filename(name) with config_path.open("w", encoding="utf-8") as f: @@ -261,9 +248,7 @@ def convert( if layout in {"monolithic", "raw_official"}: # TODO: replace model.safetensors with the official monolithic file name. - components = split_monolithic( - load_checkpoint(default_monolithic_checkpoint(src_path)) - ) + components = split_monolithic(load_checkpoint(default_monolithic_checkpoint(src_path))) elif layout in {"separate_components", "mixed"}: if not src_path.is_dir(): raise ValueError(f"{layout} layout requires a source directory: {src_path}") @@ -271,9 +256,7 @@ def convert( else: raise ValueError(f"Unsupported template layout: {layout}") - copied = ( - copy_passthrough(src_path, dst_dir) if src_path.is_dir() else [] - ) + copied = (copy_passthrough(src_path, dst_dir) if src_path.is_dir() else []) configs = build_component_configs(src_path if src_path.is_dir() else src_path.parent) validate_component_configs(configs) for name, state in components.items(): @@ -289,9 +272,7 @@ def convert( def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--src", required=True, help="HF repo id, local dir, or checkpoint path" - ) + parser.add_argument("--src", required=True, help="HF repo id, local dir, or checkpoint path") parser.add_argument("--revision", help="HF branch, tag, or commit for repo sources") parser.add_argument( "--dst", diff --git a/.agents/skills/add-model-08-trace/templates/block_trace_debug.py b/.agents/skills/add-model-08-trace/templates/block_trace_debug.py index d984be24d..45e113953 100644 --- a/.agents/skills/add-model-08-trace/templates/block_trace_debug.py +++ b/.agents/skills/add-model-08-trace/templates/block_trace_debug.py @@ -24,8 +24,8 @@ from typing import Any import torch -FAMILY: str = "" # e.g. "magi_human", "ltx2", "wan" -COMPONENT: str = "" # e.g. "dit", "vae", "encoder" +FAMILY: str = "" # e.g. "magi_human", "ltx2", "wan" +COMPONENT: str = "" # e.g. "dit", "vae", "encoder" DRILL_LAYER_ENV: str = "_DEBUG_DRILL_LAYER" HYPOTHESIS_ENV: str = "_DEBUG_PATCH_" REL_THRESHOLD: float = 0.005 # 0.5% abs_mean drift flags a block as divergent @@ -94,6 +94,7 @@ def _attach_block_hooks( handles: list[Any] = [] def _hook(name: str): + def fn(_module, _inputs, outputs): t = outputs[0] if isinstance(outputs, tuple) else outputs if not torch.is_tensor(t): @@ -101,6 +102,7 @@ def _attach_block_hooks( log.append({"side": label, **_stat(name, t)}) if tensors is not None: tensors[name] = t.detach().float().cpu() + return fn def _pre_hook(name: str): @@ -114,6 +116,7 @@ def _attach_block_hooks( log.append({"side": label, **_stat(key, t)}) if tensors is not None: tensors[key] = t.detach().float().cpu() + return fn # TODO: adapt attribute paths to your model. Remove adapter block if absent. @@ -131,43 +134,21 @@ def _attach_block_hooks( # magi-human uses: attention, mlp.pre_norm, mlp.up_gate_proj, # mlp.down_proj (pre+post), mlp, attn_post_norm, mlp_post_norm. if hasattr(layer, "attention"): - handles.append( - layer.attention.register_forward_hook(_hook(f"{tag}.attention")) - ) + handles.append(layer.attention.register_forward_hook(_hook(f"{tag}.attention"))) if hasattr(layer, "mlp"): mlp = layer.mlp if hasattr(mlp, "pre_norm"): - handles.append( - mlp.pre_norm.register_forward_hook(_hook(f"{tag}.mlp.pre_norm")) - ) + handles.append(mlp.pre_norm.register_forward_hook(_hook(f"{tag}.mlp.pre_norm"))) if hasattr(mlp, "up_gate_proj"): - handles.append( - mlp.up_gate_proj.register_forward_hook( - _hook(f"{tag}.mlp.up_gate_proj") - ) - ) + handles.append(mlp.up_gate_proj.register_forward_hook(_hook(f"{tag}.mlp.up_gate_proj"))) if hasattr(mlp, "down_proj"): - handles.append( - mlp.down_proj.register_forward_pre_hook( - _pre_hook(f"{tag}.mlp.down_proj") - ) - ) - handles.append( - mlp.down_proj.register_forward_hook(_hook(f"{tag}.mlp.down_proj")) - ) + handles.append(mlp.down_proj.register_forward_pre_hook(_pre_hook(f"{tag}.mlp.down_proj"))) + handles.append(mlp.down_proj.register_forward_hook(_hook(f"{tag}.mlp.down_proj"))) handles.append(mlp.register_forward_hook(_hook(f"{tag}.mlp"))) if hasattr(layer, "attn_post_norm"): - handles.append( - layer.attn_post_norm.register_forward_hook( - _hook(f"{tag}.attn_post_norm") - ) - ) + handles.append(layer.attn_post_norm.register_forward_hook(_hook(f"{tag}.attn_post_norm"))) if hasattr(layer, "mlp_post_norm"): - handles.append( - layer.mlp_post_norm.register_forward_hook( - _hook(f"{tag}.mlp_post_norm") - ) - ) + handles.append(layer.mlp_post_norm.register_forward_hook(_hook(f"{tag}.mlp_post_norm"))) return handles @@ -193,11 +174,9 @@ def _write_log(entries: list[dict], path: Path) -> None: path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w") as f: for e in entries: - f.write( - f"{e['name']} {e['shape']} " - f"{e['abs_mean']:.8f} {e['sum']:.4f} " - f"{e['min']:.6f} {e['max']:.6f}\n" - ) + f.write(f"{e['name']} {e['shape']} " + f"{e['abs_mean']:.8f} {e['sum']:.4f} " + f"{e['min']:.6f} {e['max']:.6f}\n") def _sort_key(name: str, drill_layer: int) -> tuple: @@ -205,9 +184,14 @@ def _sort_key(name: str, drill_layer: int) -> tuple: return (0, "") if name.startswith(f"L{drill_layer:02d}."): sub_order = { - "attention": 0, "attn_post_norm": 1, "mlp.pre_norm": 2, - "mlp.up_gate_proj": 3, "mlp.down_proj": 4, - "mlp.down_proj": 5, "mlp": 6, "mlp_post_norm": 7, + "attention": 0, + "attn_post_norm": 1, + "mlp.pre_norm": 2, + "mlp.up_gate_proj": 3, + "mlp.down_proj": 4, + "mlp.down_proj": 5, + "mlp": 6, + "mlp_post_norm": 7, }.get(name.split(".", 1)[1], 9) return (1, f"block[{drill_layer:02d}]", sub_order) if name.startswith("block["): @@ -216,10 +200,8 @@ def _sort_key(name: str, drill_layer: int) -> tuple: def _print_table(by_name: dict[str, dict], drill_layer: int) -> int | None: - hdr = ( - f"{'name':<18} {'up_shape':<22} {'up_absmean':>12} {'fv_absmean':>12} " - f"{'absmean_diff':>14} {'rel%':>8} {'up_sum':>14} {'fv_sum':>14} {'sum_diff':>12}" - ) + hdr = (f"{'name':<18} {'up_shape':<22} {'up_absmean':>12} {'fv_absmean':>12} " + f"{'absmean_diff':>14} {'rel%':>8} {'up_sum':>14} {'fv_sum':>14} {'sum_diff':>12}") print(f"\n{hdr}\n{'-' * len(hdr)}") first_div: int | None = None for name in sorted(by_name.keys(), key=lambda n: _sort_key(n, drill_layer)): @@ -235,11 +217,9 @@ def _print_table(by_name: dict[str, dict], drill_layer: int) -> int | None: flag = " <<< DIVERGE" if first_div is None: first_div = int(name[len("block["):-1]) - print( - f"{name:<18} {str(up['shape']):<22} {up['abs_mean']:>12.6f} " - f"{fv['abs_mean']:>12.6f} {am_diff:>14.6f} {am_rel * 100:>7.3f}% " - f"{up['sum']:>14.4f} {fv['sum']:>14.4f} {sum_diff:>12.4f}{flag}" - ) + print(f"{name:<18} {str(up['shape']):<22} {up['abs_mean']:>12.6f} " + f"{fv['abs_mean']:>12.6f} {am_diff:>14.6f} {am_rel * 100:>7.3f}% " + f"{up['sum']:>14.4f} {fv['sum']:>14.4f} {sum_diff:>12.4f}{flag}") return first_div @@ -255,10 +235,8 @@ def _print_elementwise(up_t: dict[str, torch.Tensor], fv_t: dict[str, torch.Tens continue diff = (a - b).abs() rel = (diff.mean().item() / max(a.abs().mean().item(), 1e-9)) * 100 - print( - f"{name:<30} {str(tuple(a.shape)):<22} " - f"{diff.max().item():>12.6f} {diff.mean().item():>12.6f} {rel:>9.4f}%" - ) + print(f"{name:<30} {str(tuple(a.shape)):<22} " + f"{diff.max().item():>12.6f} {diff.mean().item():>12.6f} {rel:>9.4f}%") def main() -> None: diff --git a/.agents/skills/add-model-09-pipeline/templates/pipeline_parity_test.py b/.agents/skills/add-model-09-pipeline/templates/pipeline_parity_test.py index 556000195..07c29c0da 100644 --- a/.agents/skills/add-model-09-pipeline/templates/pipeline_parity_test.py +++ b/.agents/skills/add-model-09-pipeline/templates/pipeline_parity_test.py @@ -43,12 +43,10 @@ def _add_official_to_path() -> Path: def _log_tensor_stats(label: str, tensor: torch.Tensor) -> None: value = tensor.detach().float() - print( - f"[{_MODEL_FAMILY} PIPELINE] {label}: shape={tuple(tensor.shape)} " - f"dtype={tensor.dtype} device={tensor.device} " - f"min={value.min().item():.6f} max={value.max().item():.6f} " - f"mean={value.mean().item():.6f} std={value.std().item():.6f}" - ) + print(f"[{_MODEL_FAMILY} PIPELINE] {label}: shape={tuple(tensor.shape)} " + f"dtype={tensor.dtype} device={tensor.device} " + f"min={value.min().item():.6f} max={value.max().item():.6f} " + f"mean={value.mean().item():.6f} std={value.std().item():.6f}") def _extract_tensor(output: Any, key: str) -> torch.Tensor: @@ -73,10 +71,8 @@ def _run_official_pipeline( device: torch.device, ) -> Any: del official_path, params, device - pytest.skip( - "TODO: import the official pipeline/factory, load official weights, " - "run with params, and return the comparison target." - ) + pytest.skip("TODO: import the official pipeline/factory, load official weights, " + "run with params, and return the comparison target.") def _run_fastvideo_pipeline(model_path: Path, params: dict[str, Any]) -> Any: @@ -146,8 +142,6 @@ def test_todo_model_family_pipeline_official_parity() -> None: assert official_tensor.shape == fastvideo_tensor.shape diff = (official_tensor - fastvideo_tensor).abs() - print( - f"diff max={diff.max().item():.6f} " - f"mean={diff.mean().item():.6f} median={diff.median().item():.6f}" - ) + print(f"diff max={diff.max().item():.6f} " + f"mean={diff.mean().item():.6f} median={diff.median().item():.6f}") assert_close(fastvideo_tensor, official_tensor, atol=1e-2, rtol=1e-2) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 91b47edd9..5acccc442 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -22,6 +22,7 @@ repos: hooks: - id: yapf args: [--in-place, --verbose] + language_version: python3.12 additional_dependencies: [toml] # TODO: Remove when yapf is upgraded - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.11.12 diff --git a/apps/dreamverse/dreamverse/tests/conftest.py b/apps/dreamverse/dreamverse/tests/conftest.py index 89b3ea01e..acc709fa4 100644 --- a/apps/dreamverse/dreamverse/tests/conftest.py +++ b/apps/dreamverse/dreamverse/tests/conftest.py @@ -3,7 +3,6 @@ from __future__ import annotations import sys from pathlib import Path - TESTS_DIR = Path(__file__).resolve().parent DREAMVERSE_PACKAGE_DIR = TESTS_DIR.parent DREAMVERSE_APP_DIR = DREAMVERSE_PACKAGE_DIR.parent diff --git a/apps/dreamverse/dreamverse/tests/test_config.py b/apps/dreamverse/dreamverse/tests/test_config.py index 99a7a0893..755bb3b40 100644 --- a/apps/dreamverse/dreamverse/tests/test_config.py +++ b/apps/dreamverse/dreamverse/tests/test_config.py @@ -5,7 +5,6 @@ from pathlib import Path import pytest - SERVER_DIR = Path(__file__).resolve().parents[1] @@ -53,9 +52,7 @@ def test_config_defaults_to_cerebras_with_parallel_groq_fallback_stage(monkeypat module = _load_config_module() assert module.PROMPT_PROVIDER == "cerebras" - assert module.PROMPT_PROVIDER_RUNTIME_STAGES == ( - ("cerebras", "groq"), - ) + assert module.PROMPT_PROVIDER_RUNTIME_STAGES == (("cerebras", "groq"), ) assert module.PROMPT_PROVIDER_PRIORITY == ( "cerebras", "groq", @@ -86,9 +83,7 @@ def test_config_ignores_legacy_groq_primary_override(monkeypatch): module = _load_config_module() assert module.PROMPT_PROVIDER == "cerebras" - assert module.PROMPT_PROVIDER_RUNTIME_STAGES == ( - ("cerebras", "groq"), - ) + assert module.PROMPT_PROVIDER_RUNTIME_STAGES == (("cerebras", "groq"), ) assert module.PROMPT_PROVIDER_PRIORITY == ( "cerebras", "groq", @@ -106,24 +101,17 @@ def test_config_uses_local_overlay_paths_when_devtools_enabled(monkeypatch, tmp_ assert module.DEVTOOLS_ENABLED is True assert module.FRONTEND_ROOT.as_posix().endswith("apps/dreamverse/web") - assert module.PROMPT_ENHANCE_SYSTEM_PROMPT_PATH.endswith( - "dreamverse/prompts.local/next_segment_system_prompt.md" - ) + assert module.PROMPT_ENHANCE_SYSTEM_PROMPT_PATH.endswith("dreamverse/prompts.local/next_segment_system_prompt.md") assert module.PROMPT_ENHANCE_SYSTEM_PROMPT_FALLBACK_PATH.endswith( - "dreamverse/prompts/next_segment_system_prompt.md" - ) + "dreamverse/prompts/next_segment_system_prompt.md") assert module.PROMPT_REWRITE_USER_SYSTEM_PROMPT_PATH.endswith( - "dreamverse/prompts.local/rewrite_user_system_prompt.md" - ) + "dreamverse/prompts.local/rewrite_user_system_prompt.md") assert module.PROMPT_REWRITE_USER_SYSTEM_PROMPT_FALLBACK_PATH.endswith( - "dreamverse/prompts/rewrite_user_system_prompt.md" - ) + "dreamverse/prompts/rewrite_user_system_prompt.md") assert module.CURATED_PRESETS_FILE_PATH.endswith( - "apps/dreamverse/web/prompts.local/selected_ltx2_continuation_story_presets.json" - ) + "apps/dreamverse/web/prompts.local/selected_ltx2_continuation_story_presets.json") assert module.CURATED_PRESETS_FALLBACK_FILE_PATH.endswith( - "apps/dreamverse/web/prompts/selected_ltx2_continuation_story_presets.json" - ) + "apps/dreamverse/web/prompts/selected_ltx2_continuation_story_presets.json") assert module.FRONTEND_STATIC_DIR_CANDIDATES[:2] == ( str(module.FRONTEND_ROOT / "out"), str(module.FRONTEND_ROOT / "dist"), diff --git a/apps/dreamverse/dreamverse/tests/test_entrypoints.py b/apps/dreamverse/dreamverse/tests/test_entrypoints.py index 15b234d25..68fa0fc59 100644 --- a/apps/dreamverse/dreamverse/tests/test_entrypoints.py +++ b/apps/dreamverse/dreamverse/tests/test_entrypoints.py @@ -9,6 +9,7 @@ from fastapi.testclient import TestClient import fastvideo.entrypoints.streaming as streaming_entrypoints import pytest + def _install_stack03_import_stubs(monkeypatch): """Keep entrypoint tests focused while later-stack runtime modules are absent.""" if not hasattr(streaming_entrypoints, "build_health_router"): @@ -17,6 +18,7 @@ def _install_stack03_import_stubs(monkeypatch): gpu_pool_stub = types.ModuleType("dreamverse.gpu_pool") class GPUPool: + def __init__(self, _gpu_ids): pass @@ -49,6 +51,7 @@ def _install_stack03_import_stubs(monkeypatch): controller_stub = types.ModuleType("dreamverse.session.controller") class SessionController: + def __init__(self, **_kwargs): pass @@ -76,13 +79,11 @@ def _run_cli(module, monkeypatch, argv: list[str]) -> list[dict[str, object]]: uvicorn_stub = types.ModuleType("uvicorn") def run(app, host: str, port: int) -> None: - calls.append( - { - "app": app, - "host": host, - "port": port, - } - ) + calls.append({ + "app": app, + "host": host, + "port": port, + }) uvicorn_stub.run = run monkeypatch.setitem(sys.modules, "uvicorn", uvicorn_stub) @@ -99,13 +100,11 @@ def test_server_cli_defaults_to_local_web_port(monkeypatch): server_main = _import_server_main(monkeypatch) calls = _run_cli(server_main, monkeypatch, ["dreamverse-server"]) - assert calls == [ - { - "app": server_main.app, - "host": "0.0.0.0", - "port": 8009, - } - ] + assert calls == [{ + "app": server_main.app, + "host": "0.0.0.0", + "port": 8009, + }] def test_server_cli_allows_explicit_host_and_port(monkeypatch): @@ -116,13 +115,11 @@ def test_server_cli_allows_explicit_host_and_port(monkeypatch): ["dreamverse-server", "--host", "127.0.0.1", "--port", "8123"], ) - assert calls == [ - { - "app": server_main.app, - "host": "127.0.0.1", - "port": 8123, - } - ] + assert calls == [{ + "app": server_main.app, + "host": "127.0.0.1", + "port": 8123, + }] def test_server_does_not_expose_backend_source_as_static_assets(monkeypatch): @@ -142,13 +139,11 @@ def test_mock_server_cli_defaults_to_local_web_port(monkeypatch): ["dreamverse-mock-server"], ) - assert calls == [ - { - "app": mock_server.app, - "host": "0.0.0.0", - "port": 8009, - } - ] + assert calls == [{ + "app": mock_server.app, + "host": "0.0.0.0", + "port": 8009, + }] def test_mock_server_cli_updates_latency(monkeypatch): @@ -161,13 +156,11 @@ def test_mock_server_cli_updates_latency(monkeypatch): ["dreamverse-mock-server", "--latency", "321", "--port", "8111"], ) - assert calls == [ - { - "app": mock_server.app, - "host": "0.0.0.0", - "port": 8111, - } - ] + assert calls == [{ + "app": mock_server.app, + "host": "0.0.0.0", + "port": 8111, + }] assert mock_server.LATENCY_MS == 321 finally: mock_server.LATENCY_MS = old_latency_ms diff --git a/apps/dreamverse/dreamverse/tests/test_gpu_pool.py b/apps/dreamverse/dreamverse/tests/test_gpu_pool.py index e30814539..e676bd14c 100644 --- a/apps/dreamverse/dreamverse/tests/test_gpu_pool.py +++ b/apps/dreamverse/dreamverse/tests/test_gpu_pool.py @@ -7,7 +7,6 @@ from types import SimpleNamespace import pytest - import dreamverse.gpu_pool as gpu_pool @@ -85,9 +84,7 @@ def test_send_command_raises_on_worker_death(): cmd_q = ctx.Queue() resp_q = ctx.Queue() - proc = ctx.Process( - target=_child_consume_and_exit, args=(cmd_q, resp_q) - ) + proc = ctx.Process(target=_child_consume_and_exit, args=(cmd_q, resp_q)) proc.start() # Wait for the spawn child to fully boot. Allow generous time — diff --git a/apps/dreamverse/dreamverse/tests/test_import_contract.py b/apps/dreamverse/dreamverse/tests/test_import_contract.py index a11312e4d..f6cee5493 100644 --- a/apps/dreamverse/dreamverse/tests/test_import_contract.py +++ b/apps/dreamverse/dreamverse/tests/test_import_contract.py @@ -7,7 +7,7 @@ ALLOWED_PREFIXES = ( "fastvideo.entrypoints.video_generator", "fastvideo.configs", ) -ALLOWED_EXACT = ("fastvideo",) +ALLOWED_EXACT = ("fastvideo", ) FORBIDDEN_PREFIXES = ( "fastvideo.pipelines", "fastvideo.models", @@ -38,19 +38,13 @@ def test_dreamverse_server_imports_only_public_fastvideo_surfaces() -> None: except SyntaxError as task_exc: raise AssertionError(f"Failed to parse {path}") from task_exc for node in ast.walk(tree): - names = ( - [a.name for a in node.names] if isinstance(node, ast.Import) - else [node.module] if isinstance(node, ast.ImportFrom) and node.module - else [] - ) + names = ([a.name for a in node.names] if isinstance(node, ast.Import) else + [node.module] if isinstance(node, ast.ImportFrom) and node.module else []) for name in names: if not name: continue rel_path = str(path.relative_to(root)) - if ( - name.startswith(FORBIDDEN_PREFIXES) - and (rel_path, name) not in ALLOWED_INTERNAL_IMPORTS - ): + if (name.startswith(FORBIDDEN_PREFIXES) and (rel_path, name) not in ALLOWED_INTERNAL_IMPORTS): bad.append((str(path.relative_to(root)), getattr(node, "lineno", 0), name)) assert bad == [], f"Forbidden internal imports: {bad}" diff --git a/apps/dreamverse/dreamverse/tests/test_mock_server.py b/apps/dreamverse/dreamverse/tests/test_mock_server.py index 30730e074..1f29665f5 100644 --- a/apps/dreamverse/dreamverse/tests/test_mock_server.py +++ b/apps/dreamverse/dreamverse/tests/test_mock_server.py @@ -6,7 +6,6 @@ import os from fastapi import WebSocketDisconnect - os.environ.setdefault("CEREBRAS_API_KEY", "dummy") os.environ.setdefault("GROQ_API_KEY", "dummy") @@ -14,6 +13,7 @@ import dreamverse.mock_server as mock_server class _FakeWebSocket: + def __init__(self, messages: list[tuple[float, dict[str, object]]]): self._messages = messages self._index = 0 @@ -49,34 +49,34 @@ def test_mock_server_matches_current_single5s_protocol(): mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes" mock_server.LATENCY_MS = 1 - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "simple_prompt_1", - "curated_prompts": ["selected prompt"], - "single_clip_mode": True, - "enhancement_enabled": False, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - ( - 0.01, - { - "type": "simple_generate", - "preset_id": "simple_custom_prompt", - "prompt_id": "simple_custom_prompt", - "prompt": "custom prompt", - "enhancement_enabled": True, - "initial_image": None, - }, - ), - (0.20, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "simple_prompt_1", + "curated_prompts": ["selected prompt"], + "single_clip_mode": True, + "enhancement_enabled": False, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "simple_generate", + "preset_id": "simple_custom_prompt", + "prompt_id": "simple_custom_prompt", + "prompt": "custom prompt", + "enhancement_enabled": True, + "initial_image": None, + }, + ), + (0.20, { + "type": "leave" + }), + ]) asyncio.run(mock_server.websocket_endpoint(ws)) @@ -92,24 +92,14 @@ def test_mock_server_matches_current_single5s_protocol(): assert message_types.count("ltx2_stream_complete") == 2 assert "prompt_sources_blocked" not in message_types - segment_start_events = [ - payload - for payload in ws.sent_json - if payload["type"] == "ltx2_segment_start" - ] - gpu_assigned_event = next( - payload for payload in ws.sent_json if payload["type"] == "gpu_assigned" - ) + segment_start_events = [payload for payload in ws.sent_json if payload["type"] == "ltx2_segment_start"] + gpu_assigned_event = next(payload for payload in ws.sent_json if payload["type"] == "gpu_assigned") assert gpu_assigned_event["session_timeout"] == mock_server.SESSION_TIMEOUT_SECONDS assert [payload["segment_idx"] for payload in segment_start_events] == [1, 1] assert segment_start_events[0]["prompt"] == "selected prompt" assert segment_start_events[1]["prompt"] == "custom prompt" - step_complete_events = [ - payload - for payload in ws.sent_json - if payload["type"] == "step_complete" - ] + step_complete_events = [payload for payload in ws.sent_json if payload["type"] == "step_complete"] assert len(step_complete_events) == 2 assert step_complete_events[0]["latency_ms"] == { "total": 121.0, @@ -134,29 +124,29 @@ def test_mock_server_regular_cap_waits_for_rewrite_rollout(): mock_server.LATENCY_MS = 1 mock_server.GENERATION_SEGMENT_CAP = 1 - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "test_preset", - "curated_prompts": ["segment one"], - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - ( - 0.02, - { - "type": "rewrite_seed_prompts", - "rewrite_instruction": "start a new rollout", - }, - ), - (0.20, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "test_preset", + "curated_prompts": ["segment one"], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.02, + { + "type": "rewrite_seed_prompts", + "rewrite_instruction": "start a new rollout", + }, + ), + (0.20, { + "type": "leave" + }), + ]) asyncio.run(mock_server.websocket_endpoint(ws)) @@ -166,11 +156,7 @@ def test_mock_server_regular_cap_waits_for_rewrite_rollout(): assert "generation_cap_reached" not in message_types assert "prompt_sources_blocked" not in message_types - segment_start_events = [ - payload - for payload in ws.sent_json - if payload["type"] == "ltx2_segment_start" - ] + segment_start_events = [payload for payload in ws.sent_json if payload["type"] == "ltx2_segment_start"] assert [payload["segment_idx"] for payload in segment_start_events] == [1, 1] assert segment_start_events[0]["prompt"] == "segment one" assert segment_start_events[1]["prompt"] == "segment one [start a new rollout]" @@ -187,54 +173,40 @@ def test_mock_server_rewrite_during_active_segment_restarts_from_first_rewritten mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes" mock_server.LATENCY_MS = 100 - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "test_preset", - "curated_prompts": ["segment one", "segment two"], - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - ( - 0.02, - { - "type": "rewrite_seed_prompts", - "rewrite_instruction": "restart from rewrite", - }, - ), - (0.40, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "test_preset", + "curated_prompts": ["segment one", "segment two"], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.02, + { + "type": "rewrite_seed_prompts", + "rewrite_instruction": "restart from rewrite", + }, + ), + (0.40, { + "type": "leave" + }), + ]) asyncio.run(mock_server.websocket_endpoint(ws)) - segment_start_events = [ - payload - for payload in ws.sent_json - if payload["type"] == "ltx2_segment_start" - ] + segment_start_events = [payload for payload in ws.sent_json if payload["type"] == "ltx2_segment_start"] assert [payload["prompt"] for payload in segment_start_events[:2]] == [ "segment one", "segment one [restart from rewrite]", ] - assert all( - payload["prompt"] != "segment two" - for payload in segment_start_events[1:] - ) - reset_events = [ - payload - for payload in ws.sent_json - if payload.get("type") == "seed_prompts_reset_applied" - ] - assert any( - payload.get("reason") == "rewrite_during_generation" - for payload in reset_events - ) + assert all(payload["prompt"] != "segment two" for payload in segment_start_events[1:]) + reset_events = [payload for payload in ws.sent_json if payload.get("type") == "seed_prompts_reset_applied"] + assert any(payload.get("reason") == "rewrite_during_generation" for payload in reset_events) finally: mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes mock_server.LATENCY_MS = old_latency_ms @@ -247,24 +219,24 @@ def test_mock_server_supports_initial_custom_rollout_prompt(): mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes" mock_server.LATENCY_MS = 1 - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "custom_editable", - "preset_label": "Custom rollout", - "curated_prompts": [], - "initial_rollout_prompt": "A moonbase corridor thriller with flooding", - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - (0.20, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "custom_editable", + "preset_label": "Custom rollout", + "curated_prompts": [], + "initial_rollout_prompt": "A moonbase corridor thriller with flooding", + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.20, { + "type": "leave" + }), + ]) asyncio.run(mock_server.websocket_endpoint(ws)) @@ -276,15 +248,9 @@ def test_mock_server_supports_initial_custom_rollout_prompt(): assert "ltx2_stream_start" in message_types assert "prompt_sources_blocked" not in message_types - segment_start_events = [ - payload - for payload in ws.sent_json - if payload["type"] == "ltx2_segment_start" - ] + segment_start_events = [payload for payload in ws.sent_json if payload["type"] == "ltx2_segment_start"] assert segment_start_events - assert segment_start_events[0]["prompt"] == ( - "A moonbase corridor thriller with flooding [segment 1]" - ) + assert segment_start_events[0]["prompt"] == ("A moonbase corridor thriller with flooding [segment 1]") finally: mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes mock_server.LATENCY_MS = old_latency_ms @@ -297,35 +263,37 @@ def test_mock_server_can_start_new_project_without_reconnecting(): mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes" mock_server.LATENCY_MS = 40 - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "test_preset", - "curated_prompts": ["segment one"], - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - (0.02, {"type": "end_project_keep_session"}), - ( - 0.20, - { - "type": "project_init_v1", - "preset_id": "test_preset_2", - "preset_label": "Test Preset 2", - "curated_prompts": ["segment two"], - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - (0.40, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "test_preset", + "curated_prompts": ["segment one"], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.02, { + "type": "end_project_keep_session" + }), + ( + 0.20, + { + "type": "project_init_v1", + "preset_id": "test_preset_2", + "preset_label": "Test Preset 2", + "curated_prompts": ["segment two"], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.40, { + "type": "leave" + }), + ]) asyncio.run(mock_server.websocket_endpoint(ws)) @@ -336,16 +304,11 @@ def test_mock_server_can_start_new_project_without_reconnecting(): project_idle_index = message_types.index("project_idle") stream_start_indexes = [ - index for index, message_type in enumerate(message_types) - if message_type == "ltx2_stream_start" + index for index, message_type in enumerate(message_types) if message_type == "ltx2_stream_start" ] assert stream_start_indexes[0] < project_idle_index < stream_start_indexes[1] - segment_start_events = [ - payload - for payload in ws.sent_json - if payload["type"] == "ltx2_segment_start" - ] + segment_start_events = [payload for payload in ws.sent_json if payload["type"] == "ltx2_segment_start"] assert [payload["prompt"] for payload in segment_start_events[:2]] == [ "segment one", "segment two", diff --git a/apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py b/apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py index 52552b16a..49bfeda8a 100644 --- a/apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py +++ b/apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py @@ -6,7 +6,6 @@ import os import re import time - os.environ.setdefault("CEREBRAS_API_KEY", "dummy") os.environ.setdefault("GROQ_API_KEY", "dummy") @@ -22,6 +21,7 @@ from dreamverse.prompt_enhancer import ( class _FakeResponse: + def __init__(self, payload: dict): self._payload = payload @@ -30,6 +30,7 @@ class _FakeResponse: class _FakeSyncCompletions: + def __init__(self, payload: dict): self._payload = payload @@ -38,6 +39,7 @@ class _FakeSyncCompletions: class _FakeSyncClient: + def __init__(self, payload: dict): self.chat = type( "_FakeChat", @@ -47,6 +49,7 @@ class _FakeSyncClient: class _DelayedSyncCompletions: + def __init__(self, payload: dict, delay_s: float = 0.0, exc: Exception | None = None): self._payload = payload self._delay_s = delay_s @@ -61,29 +64,26 @@ class _DelayedSyncCompletions: class _DelayedSyncClient: + def __init__(self, payload: dict, delay_s: float = 0.0, exc: Exception | None = None): self.chat = type( "_FakeChat", (), - { - "completions": _DelayedSyncCompletions( - payload, - delay_s=delay_s, - exc=exc, - ) - }, + {"completions": _DelayedSyncCompletions( + payload, + delay_s=delay_s, + exc=exc, + )}, )() def _chat_payload_with_content(content: str) -> dict: return { - "choices": [ - { - "message": { - "content": content, - } + "choices": [{ + "message": { + "content": content, } - ] + }] } @@ -172,6 +172,7 @@ def _build_staged_enhancer( class _FakeOpenAIClient: + def __init__(self, **kwargs): self.kwargs = kwargs self.chat = type( @@ -182,6 +183,7 @@ class _FakeOpenAIClient: class _FakeCerebrasClient: + def __init__(self, **kwargs): self.kwargs = kwargs self.chat = type( @@ -192,16 +194,12 @@ class _FakeCerebrasClient: def test_parse_json_response_accepts_fenced_json_with_prose(): - parsed = _parse_json_response( - "Here is the rewrite:\n```json\n{\"segment_prompts\":[\"A\",\"B\"]}\n```\nThanks." - ) + parsed = _parse_json_response("Here is the rewrite:\n```json\n{\"segment_prompts\":[\"A\",\"B\"]}\n```\nThanks.") assert parsed == {"segment_prompts": ["A", "B"]} def test_parse_json_response_extracts_first_embedded_object(): - parsed = _parse_json_response( - "Model output:\n{\"segment_prompts\":[\"A\",\"B\"]}\n(complete)" - ) + parsed = _parse_json_response("Model output:\n{\"segment_prompts\":[\"A\",\"B\"]}\n(complete)") assert parsed == {"segment_prompts": ["A", "B"]} @@ -268,16 +266,12 @@ def test_build_client_supports_groq_provider(monkeypatch): def test_rewrite_prompt_sequence_accepts_segment_prompts_output(): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) result = asyncio.run( enhancer.rewrite_prompt_sequence( ["prompt one", "prompt two"], rewrite_instruction="make it cinematic", - ) - ) + )) assert result.fallback_used is False assert result.error is None assert result.rollout_id == "preset_a" @@ -286,15 +280,12 @@ def test_rewrite_prompt_sequence_accepts_segment_prompts_output(): def test_rewrite_prompt_sequence_accepts_legacy_rewritten_prompts_output(): - enhancer = _build_test_enhancer( - _chat_payload_with_content('{"rewritten_prompts":["A","B"]}') - ) + enhancer = _build_test_enhancer(_chat_payload_with_content('{"rewritten_prompts":["A","B"]}')) result = asyncio.run( enhancer.rewrite_prompt_sequence( ["prompt one", "prompt two"], rewrite_instruction="make it cinematic", - ) - ) + )) assert result.fallback_used is False assert result.error is None assert result.rollout_id == "current_rollout" @@ -303,19 +294,14 @@ def test_rewrite_prompt_sequence_accepts_legacy_rewritten_prompts_output(): def test_rewrite_prompt_sequence_accepts_segment_dicts_without_top_level_rollout_metadata(): - enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"segments":[{"prompt":"A"},{"text":"B"}]}' - ) - ) + enhancer = _build_test_enhancer(_chat_payload_with_content('{"segments":[{"prompt":"A"},{"text":"B"}]}')) result = asyncio.run( enhancer.rewrite_prompt_sequence( ["prompt one", "prompt two"], preset_id="preset_a", preset_label="Preset A", rewrite_instruction="make it cinematic", - ) - ) + )) assert result.fallback_used is False assert result.error is None assert result.rollout_id == "preset_a" @@ -329,14 +315,12 @@ def test_rewrite_prompt_sequence_accepts_numbered_prose_output(): "The user is asking for a cinematic rewrite.\n\n" "1. A dog bounds across the moon's dusty surface, kicking up silver regolith as it chases a rabbit beneath the black sky.\n" "2. The rabbit darts around a crater rim while the dog lunges after it, Earth glowing blue in the distance.\n" - ) - ) + )) result = asyncio.run( enhancer.rewrite_prompt_sequence( ["prompt one", "prompt two"], rewrite_instruction="make it cinematic", - ) - ) + )) assert result.fallback_used is False assert result.error is None assert result.rollout_id == "current_rollout" @@ -355,12 +339,10 @@ def test_enhance_prompt_prefers_cerebras_before_groq_fallback(): groq_delay_s=0.01, ) - result = asyncio.run( - enhancer.enhance_prompt( - "A rainy alley at night", - mode="single_clip", - ) - ) + result = asyncio.run(enhancer.enhance_prompt( + "A rainy alley at night", + mode="single_clip", + )) assert result.fallback_used is False assert result.error is None @@ -381,12 +363,10 @@ def test_enhance_prompt_uses_groq_when_cerebras_fails(): groq_delay_s=0.01, ) - result = asyncio.run( - enhancer.enhance_prompt( - "A rainy alley at night", - mode="single_clip", - ) - ) + result = asyncio.run(enhancer.enhance_prompt( + "A rainy alley at night", + mode="single_clip", + )) assert result.fallback_used is False assert result.error is None @@ -408,13 +388,11 @@ def test_enhance_prompt_can_use_groq_when_cerebras_times_out(): enhancer.http_timeout_ms = 50 enhancer.default_timeout_ms = 50 - result = asyncio.run( - enhancer.enhance_prompt( - "A rainy alley at night", - mode="single_clip", - timeout_ms=50, - ) - ) + result = asyncio.run(enhancer.enhance_prompt( + "A rainy alley at night", + mode="single_clip", + timeout_ms=50, + )) assert result.fallback_used is False assert result.error is None @@ -434,12 +412,10 @@ def test_enhance_prompt_can_use_cerebras_when_it_returns_first(): groq_delay_s=0.08, ) - result = asyncio.run( - enhancer.enhance_prompt( - "A rainy alley at night", - mode="single_clip", - ) - ) + result = asyncio.run(enhancer.enhance_prompt( + "A rainy alley at night", + mode="single_clip", + )) assert result.fallback_used is False assert result.error is None @@ -453,15 +429,12 @@ def test_enhance_prompt_can_use_cerebras_when_it_returns_first(): def test_rewrite_prompt_sequence_keeps_raw_output_on_parse_error(): - enhancer = _build_test_enhancer( - _chat_payload_with_content("I cannot comply with JSON right now.") - ) + enhancer = _build_test_enhancer(_chat_payload_with_content("I cannot comply with JSON right now.")) result = asyncio.run( enhancer.rewrite_prompt_sequence( ["prompt one", "prompt two"], rewrite_instruction="make it cinematic", - ) - ) + )) assert result.fallback_used is True assert "No JSON object found in assistant response." in (result.error or "") assert result.raw_response_text == "I cannot comply with JSON right now." @@ -473,9 +446,7 @@ def test_rewrite_prompt_sequence_keeps_raw_output_on_parse_error(): def test_rewrite_prompt_sequence_uses_current_rollout_payload_shape(): enhancer = _build_test_enhancer( _chat_payload_with_content( - '{"id":"rewritten_rollout","label":"Rewritten Rollout","segment_prompts":["A","B"]}' - ) - ) + '{"id":"rewritten_rollout","label":"Rewritten Rollout","segment_prompts":["A","B"]}')) captured = { "body": None, "timeout_seconds": None, @@ -486,8 +457,7 @@ def test_rewrite_prompt_sequence_uses_current_rollout_payload_shape(): captured["timeout_seconds"] = timeout_seconds return ( _chat_payload_with_content( - '{"id":"rewritten_rollout","label":"Rewritten Rollout","segment_prompts":["A","B"]}' - ), + '{"id":"rewritten_rollout","label":"Rewritten Rollout","segment_prompts":["A","B"]}'), '{"id":"rewritten_rollout","label":"Rewritten Rollout","segment_prompts":["A","B"]}', ) @@ -502,8 +472,7 @@ def test_rewrite_prompt_sequence_uses_current_rollout_payload_shape(): rewrite_model="gpt-test", rewrite_temperature=0.2, timeout_ms=800, - ) - ) + )) assert result.fallback_used is False assert captured["body"]["messages"][0] == { @@ -512,12 +481,12 @@ def test_rewrite_prompt_sequence_uses_current_rollout_payload_shape(): } assert captured["body"]["messages"][1]["role"] == "user" assert prompt_enhancer_module.json.loads(captured["body"]["messages"][1]["content"]) == { - "mode": "edit_existing_rollout", - "request": ( - "Rewrite all segment prompts with improved continuity and cinematic detail. " - "Keep count and ordering identical." - ), - "user_instruction": "make it cinematic", + "mode": + "edit_existing_rollout", + "request": ("Rewrite all segment prompts with improved continuity and cinematic detail. " + "Keep count and ordering identical."), + "user_instruction": + "make it cinematic", "current_rollout": { "id": "preset_a", "label": "Preset A", @@ -528,11 +497,8 @@ def test_rewrite_prompt_sequence_uses_current_rollout_payload_shape(): def test_rewrite_prompt_sequence_supports_new_rollout_mode(): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"custom_editable","label":"Custom rollout","segment_prompts":[' - '"A","B","C","D","E","F"]}' - ) - ) + _chat_payload_with_content('{"id":"custom_editable","label":"Custom rollout","segment_prompts":[' + '"A","B","C","D","E","F"]}')) captured = { "body": None, } @@ -541,10 +507,8 @@ def test_rewrite_prompt_sequence_supports_new_rollout_mode(): del timeout_seconds captured["body"] = body return ( - _chat_payload_with_content( - '{"id":"custom_editable","label":"Custom rollout","segment_prompts":[' - '"A","B","C","D","E","F"]}' - ), + _chat_payload_with_content('{"id":"custom_editable","label":"Custom rollout","segment_prompts":[' + '"A","B","C","D","E","F"]}'), '{"id":"custom_editable","label":"Custom rollout","segment_prompts":[' '"A","B","C","D","E","F"]}', ) @@ -560,30 +524,29 @@ def test_rewrite_prompt_sequence_supports_new_rollout_mode(): rewrite_model="gpt-test", rewrite_temperature=0.2, timeout_ms=800, - ) - ) + )) assert result.fallback_used is False assert result.prompts == ["A", "B", "C", "D", "E", "F"] assert prompt_enhancer_module.json.loads(captured["body"]["messages"][1]["content"]) == { - "mode": "new_rollout", - "request": ( - "Rewrite all segment prompts with improved continuity and cinematic detail. " - "Keep count and ordering identical." - ), - "user_instruction": "A moonbase corridor thriller with flooding and red alarms", - "desired_segment_count": 6, - "rollout_id_hint": "custom_editable", - "rollout_label_hint": "Custom rollout", + "mode": + "new_rollout", + "request": ("Rewrite all segment prompts with improved continuity and cinematic detail. " + "Keep count and ordering identical."), + "user_instruction": + "A moonbase corridor thriller with flooding and red alarms", + "desired_segment_count": + 6, + "rollout_id_hint": + "custom_editable", + "rollout_label_hint": + "Custom rollout", } def test_rewrite_prompt_sequence_uses_session_override_system_prompt(): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) enhancer.rewrite_all_system_prompt = "shared system prompt" captured = { "body": None, @@ -593,9 +556,7 @@ def test_rewrite_prompt_sequence_uses_session_override_system_prompt(): del timeout_seconds captured["body"] = body return ( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ), + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}'), '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}', ) @@ -609,8 +570,7 @@ def test_rewrite_prompt_sequence_uses_session_override_system_prompt(): rewrite_instruction="make it cinematic", rewrite_model="gpt-test", system_prompt_override="session specific system prompt", - ) - ) + )) assert result.fallback_used is False assert captured["body"]["messages"][0] == { @@ -621,10 +581,7 @@ def test_rewrite_prompt_sequence_uses_session_override_system_prompt(): def test_resolve_rewrite_new_rollout_system_prompt_uses_dedicated_prompt(): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) enhancer.rewrite_all_system_prompt = "shared rewrite system prompt" enhancer.rewrite_user_system_prompt = "new rollout rewrite system prompt" @@ -635,24 +592,17 @@ def test_resolve_rewrite_new_rollout_system_prompt_uses_dedicated_prompt(): def test_resolve_rewrite_new_rollout_system_prompt_prefers_override(): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) enhancer.rewrite_all_system_prompt = "shared rewrite system prompt" enhancer.rewrite_user_system_prompt = "new rollout rewrite system prompt" - resolved = enhancer.resolve_rewrite_new_rollout_system_prompt( - "session specific system prompt" - ) + resolved = enhancer.resolve_rewrite_new_rollout_system_prompt("session specific system prompt") assert resolved == "session specific system prompt" def test_generate_auto_prompt_uses_selected_model(): - enhancer = _build_test_enhancer( - _chat_payload_with_content('{"next_prompt":"Auto next"}') - ) + enhancer = _build_test_enhancer(_chat_payload_with_content('{"next_prompt":"Auto next"}')) enhancer.auto_system_prompt = "auto system prompt" enhancer.rewrite_model_options = ["gpt-test", "gpt-alt"] enhancer.rewrite_default_model = "gpt-test" @@ -680,8 +630,7 @@ def test_generate_auto_prompt_uses_selected_model(): next_segment_idx=2, model="gpt-alt", timeout_ms=800, - ) - ) + )) assert result.fallback_used is False assert result.error is None assert result.prompt == "Auto next" @@ -690,9 +639,7 @@ def test_generate_auto_prompt_uses_selected_model(): def test_enhance_prompt_uses_selected_model(): - enhancer = _build_test_enhancer( - _chat_payload_with_content('{"next_prompt":"Enhanced next"}') - ) + enhancer = _build_test_enhancer(_chat_payload_with_content('{"next_prompt":"Enhanced next"}')) enhancer.enhance_system_prompt = "enhance system prompt" enhancer.auto_system_prompt = "auto system prompt" enhancer.rewrite_model_options = ["gpt-test", "gpt-alt"] @@ -722,8 +669,7 @@ def test_enhance_prompt_uses_selected_model(): next_segment_idx=2, model="gpt-alt", timeout_ms=800, - ) - ) + )) assert result.fallback_used is False assert result.error is None assert result.prompt == "Enhanced next" @@ -732,9 +678,7 @@ def test_enhance_prompt_uses_selected_model(): def test_enhance_prompt_single_clip_uses_auto_extension_prompt_and_prompt_field(): - enhancer = _build_test_enhancer( - _chat_payload_with_content('{"prompt":"Extended single clip"}') - ) + enhancer = _build_test_enhancer(_chat_payload_with_content('{"prompt":"Extended single clip"}')) enhancer.enhance_system_prompt = "enhance system prompt" enhancer.auto_system_prompt = "auto system prompt" enhancer.rewrite_model_options = ["gpt-test", "gpt-alt"] @@ -764,14 +708,12 @@ def test_enhance_prompt_single_clip_uses_auto_extension_prompt_and_prompt_field( enhancer._request_content = _fake_request_content # type: ignore[attr-defined] - result = asyncio.run( - enhancer.enhance_prompt( - "short 5s idea", - mode="single_clip", - model="gpt-alt", - timeout_ms=800, - ) - ) + result = asyncio.run(enhancer.enhance_prompt( + "short 5s idea", + mode="single_clip", + model="gpt-alt", + timeout_ms=800, + )) assert result.fallback_used is False assert result.error is None assert result.prompt == "Extended single clip" @@ -784,17 +726,15 @@ def test_enhance_prompt_single_clip_uses_auto_extension_prompt_and_prompt_field( "single 5-second LTX-2.3 video clip. Respond with " 'valid JSON only as {"prompt": "..."}.' # noqa: E501 ), - "user_prompt": "short 5s idea", + "user_prompt": + "short 5s idea", } def test_enhance_prompt_single_clip_rejects_plain_text_response(): enhancer = _build_test_enhancer( - _chat_payload_with_content( - "Medium shot of a woman by a rainy cafe window as she lifts her " - "phone, exhales softly, and the camera makes a slow push in." - ) - ) + _chat_payload_with_content("Medium shot of a woman by a rainy cafe window as she lifts her " + "phone, exhales softly, and the camera makes a slow push in.")) enhancer.auto_system_prompt = "auto system prompt" result = asyncio.run( @@ -803,17 +743,14 @@ def test_enhance_prompt_single_clip_rejects_plain_text_response(): mode="single_clip", model="gpt-test", timeout_ms=800, - ) - ) + )) assert result.fallback_used is True assert "No JSON object found in assistant response." in result.error assert result.prompt == "" def test_enhance_prompt_single_clip_rejects_segment_prompts_json(): - enhancer = _build_test_enhancer( - _chat_payload_with_content('{"segment_prompts":["A","B"]}') - ) + enhancer = _build_test_enhancer(_chat_payload_with_content('{"segment_prompts":["A","B"]}')) enhancer.auto_system_prompt = "auto system prompt" enhancer.rewrite_model_options = ["gpt-test"] enhancer.rewrite_default_model = "gpt-test" @@ -824,17 +761,14 @@ def test_enhance_prompt_single_clip_rejects_segment_prompts_json(): mode="single_clip", model="gpt-test", timeout_ms=800, - ) - ) + )) assert result.fallback_used is True assert result.prompt == "" assert "Missing prompt string." in (result.error or "") def test_enhance_prompt_requires_json_and_does_not_fallback_to_raw_text(): - enhancer = _build_test_enhancer( - _chat_payload_with_content("A cinematic continuation with slow dolly movement.") - ) + enhancer = _build_test_enhancer(_chat_payload_with_content("A cinematic continuation with slow dolly movement.")) enhancer.enhance_system_prompt = "enhance system prompt" enhancer.rewrite_model_options = ["gpt-test"] enhancer.rewrite_default_model = "gpt-test" @@ -846,17 +780,14 @@ def test_enhance_prompt_requires_json_and_does_not_fallback_to_raw_text(): next_segment_idx=2, model="gpt-test", timeout_ms=800, - ) - ) + )) assert result.fallback_used is True assert result.prompt == "" assert "No JSON object found in assistant response." in (result.error or "") def test_generate_auto_prompt_requires_json_and_does_not_fallback_to_raw_text(): - enhancer = _build_test_enhancer( - _chat_payload_with_content("A calm, grounded continuation with subtle motion.") - ) + enhancer = _build_test_enhancer(_chat_payload_with_content("A calm, grounded continuation with subtle motion.")) enhancer.auto_system_prompt = "auto system prompt" enhancer.rewrite_model_options = ["gpt-test"] enhancer.rewrite_default_model = "gpt-test" @@ -867,34 +798,30 @@ def test_generate_auto_prompt_requires_json_and_does_not_fallback_to_raw_text(): next_segment_idx=2, model="gpt-test", timeout_ms=800, - ) - ) + )) assert result.fallback_used is True assert result.prompt == "" assert "No JSON object found in assistant response." in (result.error or "") def test_rewrite_prompt_sequence_includes_raw_json_when_content_empty(): - enhancer = _build_test_enhancer( - { - "choices": [ - { - "finish_reason": "length", - "message": { - "content": [], - "refusal": None, - }, - } - ], - "usage": {"completion_tokens": 0}, - } - ) + enhancer = _build_test_enhancer({ + "choices": [{ + "finish_reason": "length", + "message": { + "content": [], + "refusal": None, + }, + }], + "usage": { + "completion_tokens": 0 + }, + }) result = asyncio.run( enhancer.rewrite_prompt_sequence( ["prompt one", "prompt two"], rewrite_instruction="make it cinematic", - ) - ) + )) assert result.fallback_used is True assert "No rewrite segment prompts found in assistant response." in (result.error or "") assert isinstance(result.raw_response_text, str) @@ -903,10 +830,7 @@ def test_rewrite_prompt_sequence_includes_raw_json_when_content_empty(): def test_get_rewrite_model_config_returns_fixed_defaults(): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) enhancer.rewrite_default_model = "gpt-oss-120b" enhancer.rewrite_model_options = ["gpt-oss-120b"] @@ -918,10 +842,7 @@ def test_get_rewrite_model_config_returns_fixed_defaults(): def test_get_prompt_config_includes_auto_extension_prompt(): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) enhancer.enhance_system_prompt_path = "/tmp/next.md" enhancer.auto_system_prompt_path = "/tmp/auto.md" enhancer.rewrite_all_system_prompt_path = "/tmp/rewrite.md" @@ -948,19 +869,14 @@ def test_get_prompt_config_includes_auto_extension_prompt(): def test_get_prompt_config_reports_loaded_fallback_prompt_path(tmp_path): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) rewrite_fallback_path = tmp_path / "rewrite_window_system_prompt.md" rewrite_fallback_path.write_text("rewrite prompt\n", encoding="utf-8") next_path = tmp_path / "next.md" next_path.write_text("next prompt\n", encoding="utf-8") auto_path = tmp_path / "auto.md" auto_path.write_text("auto prompt\n", encoding="utf-8") - enhancer.rewrite_all_system_prompt_path = str( - tmp_path / "prompts.local" / "rewrite_window_system_prompt.md" - ) + enhancer.rewrite_all_system_prompt_path = str(tmp_path / "prompts.local" / "rewrite_window_system_prompt.md") enhancer.rewrite_all_system_prompt_fallback_path = str(rewrite_fallback_path) enhancer.enhance_system_prompt_path = str(next_path) enhancer.auto_system_prompt_path = str(auto_path) @@ -973,14 +889,9 @@ def test_get_prompt_config_reports_loaded_fallback_prompt_path(tmp_path): assert config["rewrite_window_system_prompt_path"] == str(rewrite_fallback_path) -def test_reload_system_prompts_falls_back_to_rewrite_window_when_user_prompt_empty( - tmp_path, -): +def test_reload_system_prompts_falls_back_to_rewrite_window_when_user_prompt_empty(tmp_path, ): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) next_path = tmp_path / "next.md" auto_path = tmp_path / "auto.md" rewrite_path = tmp_path / "rewrite_window_system_prompt.md" @@ -1007,10 +918,7 @@ def test_reload_system_prompts_falls_back_to_rewrite_window_when_user_prompt_emp def test_save_prompt_config_updates_auto_extension_prompt(tmp_path): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) next_path = tmp_path / "next.md" auto_path = tmp_path / "auto.md" rewrite_path = tmp_path / "rewrite.md" @@ -1021,9 +929,7 @@ def test_save_prompt_config_updates_auto_extension_prompt(tmp_path): enhancer.auto_system_prompt_path = str(auto_path) enhancer.rewrite_all_system_prompt_path = str(rewrite_path) - config = enhancer.save_prompt_config( - auto_extension_system_prompt="auto updated", - ) + config = enhancer.save_prompt_config(auto_extension_system_prompt="auto updated", ) assert auto_path.read_text(encoding="utf-8").strip() == "auto updated" assert config["auto_extension_system_prompt"] == "auto updated" @@ -1031,10 +937,7 @@ def test_save_prompt_config_updates_auto_extension_prompt(tmp_path): def test_save_prompt_config_updates_rewrite_user_prompt(tmp_path): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) next_path = tmp_path / "next.md" auto_path = tmp_path / "auto.md" rewrite_path = tmp_path / "rewrite.md" @@ -1052,9 +955,7 @@ def test_save_prompt_config_updates_rewrite_user_prompt(tmp_path): enhancer.rewrite_all_system_prompt_fallback_path = None enhancer.rewrite_user_system_prompt_fallback_path = None - config = enhancer.save_prompt_config( - rewrite_user_system_prompt="rewrite user updated", - ) + config = enhancer.save_prompt_config(rewrite_user_system_prompt="rewrite user updated", ) assert rewrite_user_path.read_text(encoding="utf-8").strip() == "rewrite user updated" assert config["rewrite_user_system_prompt"] == "rewrite user updated" @@ -1062,10 +963,7 @@ def test_save_prompt_config_updates_rewrite_user_prompt(tmp_path): def test_save_prompt_config_updates_rewrite_model(tmp_path): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) next_path = tmp_path / "next.md" auto_path = tmp_path / "auto.md" rewrite_path = tmp_path / "rewrite.md" @@ -1081,9 +979,7 @@ def test_save_prompt_config_updates_rewrite_model(tmp_path): enhancer.rewrite_default_model = "gpt-test" enhancer.rewrite_model_options = ["gpt-test", "gpt-alt"] - config = enhancer.save_prompt_config( - rewrite_model="gpt-alt", - ) + config = enhancer.save_prompt_config(rewrite_model="gpt-alt", ) assert enhancer.rewrite_default_model == "gpt-alt" assert config["rewrite_model"] == "gpt-alt" @@ -1092,10 +988,7 @@ def test_save_prompt_config_updates_rewrite_model(tmp_path): def test_save_prompt_config_updates_rewrite_temperature(tmp_path): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) next_path = tmp_path / "next.md" auto_path = tmp_path / "auto.md" rewrite_path = tmp_path / "rewrite.md" @@ -1109,9 +1002,7 @@ def test_save_prompt_config_updates_rewrite_temperature(tmp_path): enhancer.auto_system_prompt_fallback_path = None enhancer.rewrite_all_system_prompt_fallback_path = None - config = enhancer.save_prompt_config( - rewrite_temperature=1.3, - ) + config = enhancer.save_prompt_config(rewrite_temperature=1.3, ) assert enhancer.rewrite_default_temperature == 1.3 assert config["rewrite_temperature"] == 1.3 @@ -1119,10 +1010,7 @@ def test_save_prompt_config_updates_rewrite_temperature(tmp_path): def test_save_prompt_config_creates_versioned_backup_for_existing_prompt(tmp_path): enhancer = _build_test_enhancer( - _chat_payload_with_content( - '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' - ) - ) + _chat_payload_with_content('{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}')) next_path = tmp_path / "next.md" auto_path = tmp_path / "auto.md" rewrite_path = tmp_path / "rewrite_window_system_prompt.md" @@ -1136,13 +1024,9 @@ def test_save_prompt_config_creates_versioned_backup_for_existing_prompt(tmp_pat enhancer.auto_system_prompt_fallback_path = None enhancer.rewrite_all_system_prompt_fallback_path = None - enhancer.save_prompt_config( - rewrite_window_system_prompt="rewrite updated", - ) + enhancer.save_prompt_config(rewrite_window_system_prompt="rewrite updated", ) - backup_paths = sorted( - tmp_path.glob("rewrite_window_system_prompt.*.bak.md") - ) + backup_paths = sorted(tmp_path.glob("rewrite_window_system_prompt.*.bak.md")) assert rewrite_path.read_text(encoding="utf-8").strip() == "rewrite updated" assert len(backup_paths) == 1 diff --git a/apps/dreamverse/dreamverse/tests/test_realtime_stress.py b/apps/dreamverse/dreamverse/tests/test_realtime_stress.py index 2c492b351..26eeefed2 100644 --- a/apps/dreamverse/dreamverse/tests/test_realtime_stress.py +++ b/apps/dreamverse/dreamverse/tests/test_realtime_stress.py @@ -27,13 +27,8 @@ try: except ModuleNotFoundError: websockets = None # type: ignore[assignment] - -DEFAULT_PRESET_FILE = ( - Path(__file__).resolve().parents[2] - / "web" - / "prompts" - / "selected_ltx2_continuation_story_presets.json" -) +DEFAULT_PRESET_FILE = (Path(__file__).resolve().parents[2] / "web" / "prompts" / + "selected_ltx2_continuation_story_presets.json") def utc_now_iso() -> str: @@ -65,10 +60,7 @@ def safe_percentile(values: list[float], percentile: float) -> float | None: if lower == upper: return sorted_values[lower] fraction = rank - lower - return ( - sorted_values[lower] - + (sorted_values[upper] - sorted_values[lower]) * fraction - ) + return (sorted_values[lower] + (sorted_values[upper] - sorted_values[lower]) * fraction) def summarize_series(values: list[float]) -> dict[str, float | int | None]: @@ -145,24 +137,16 @@ def load_curated_prompts( selected_id = str(selected.get("id", "")).strip() or "unknown_preset" raw_prompts = selected.get("segment_prompts", []) if not isinstance(raw_prompts, list): - raise ValueError( - f"Preset {selected_id} has invalid segment_prompts (must be list)." - ) + raise ValueError(f"Preset {selected_id} has invalid segment_prompts (must be list).") - prompts = [ - str(prompt).strip() - for prompt in raw_prompts - if isinstance(prompt, str) and str(prompt).strip() - ] + prompts = [str(prompt).strip() for prompt in raw_prompts if isinstance(prompt, str) and str(prompt).strip()] if not prompts: raise ValueError(f"Preset {selected_id} has no non-empty prompts.") limited = prompts[:curated_limit] if not limited: - raise ValueError( - f"curated_limit={curated_limit} produced no prompts for preset " - f"{selected_id}." - ) + raise ValueError(f"curated_limit={curated_limit} produced no prompts for preset " + f"{selected_id}.") return selected_id, limited, len(prompts) @@ -224,11 +208,11 @@ async def run_single_session( try: async with websockets.connect( - url, - max_size=None, - ping_interval=None, - open_timeout=connect_timeout_s, - close_timeout=2.0, + url, + max_size=None, + ping_interval=None, + open_timeout=connect_timeout_s, + close_timeout=2.0, ) as ws: connect_finish_monotonic = time.monotonic() session_data["connect_finish_ts_utc"] = utc_now_iso() @@ -249,9 +233,7 @@ async def run_single_session( timeout_remaining = session_timeout_s - elapsed_s if timeout_remaining <= 0: session_data["status"] = "timeout" - session_data["error"] = ( - f"Session timed out after {session_timeout_s:.1f}s." - ) + session_data["error"] = (f"Session timed out after {session_timeout_s:.1f}s.") break recv_start_epoch = time.time() @@ -265,9 +247,7 @@ async def run_single_session( ) except asyncio.TimeoutError: session_data["status"] = "timeout" - session_data["error"] = ( - "Timed out waiting for websocket message." - ) + session_data["error"] = ("Timed out waiting for websocket message.") break except Exception as exc: session_data["status"] = "failed" @@ -288,20 +268,16 @@ async def run_single_session( chunk_gap_ms: float | None = None if last_chunk_finish_monotonic is not None: - chunk_gap_ms = ( - recv_finish_monotonic - last_chunk_finish_monotonic - ) * 1000.0 + chunk_gap_ms = (recv_finish_monotonic - last_chunk_finish_monotonic) * 1000.0 - session_data["chunks"].append( - { - "segment_idx": current_segment_idx, - "chunk_idx": session_data["total_chunks"], - "size_bytes": len(message), - "chunk_start_ts_utc": recv_start_iso, - "chunk_finish_ts_utc": recv_finish_iso, - "chunk_gap_ms": chunk_gap_ms, - } - ) + session_data["chunks"].append({ + "segment_idx": current_segment_idx, + "chunk_idx": session_data["total_chunks"], + "size_bytes": len(message), + "chunk_start_ts_utc": recv_start_iso, + "chunk_finish_ts_utc": recv_finish_iso, + "chunk_gap_ms": chunk_gap_ms, + }) last_chunk_finish_monotonic = recv_finish_monotonic last_chunk_finish_epoch = recv_finish_epoch session_data["last_chunk_finish_ts_utc"] = recv_finish_iso @@ -321,9 +297,7 @@ async def run_single_session( if msg_type == "gpu_assigned": session_data["gpu_assigned_ts_utc"] = recv_finish_iso if connect_finish_monotonic is not None: - session_data["queue_wait_ms"] = ( - recv_finish_monotonic - connect_finish_monotonic - ) * 1000.0 + session_data["queue_wait_ms"] = (recv_finish_monotonic - connect_finish_monotonic) * 1000.0 elif msg_type == "ltx2_stream_start": if initial_total_segments is None: parsed_total = parse_int(data.get("total_segments")) @@ -338,20 +312,13 @@ async def run_single_session( session_data["media_segments_completed"] += 1 if first_media_segment_complete_epoch is None: first_media_segment_complete_epoch = recv_finish_epoch - session_data[ - "first_media_segment_complete_ts_utc" - ] = recv_finish_iso + session_data["first_media_segment_complete_ts_utc"] = recv_finish_iso elif msg_type == "ltx2_segment_complete": session_data["segments_completed"] += 1 seg_idx = parse_int(data.get("segment_idx")) - if ( - initial_total_segments is not None - and seg_idx is not None - and seg_idx >= initial_total_segments - ): - session_data[ - "target_segment_complete_ts_utc" - ] = recv_finish_iso + if (initial_total_segments is not None and seg_idx is not None + and seg_idx >= initial_total_segments): + session_data["target_segment_complete_ts_utc"] = recv_finish_iso await asyncio.sleep(post_complete_wait_s) session_data["leave_sent_ts_utc"] = utc_now_iso() try: @@ -362,15 +329,11 @@ async def run_single_session( break elif msg_type == "session_timeout": session_data["status"] = "timeout" - session_data["error"] = str( - data.get("message") or "Backend session timeout" - ) + session_data["error"] = str(data.get("message") or "Backend session timeout") break elif msg_type == "error": session_data["status"] = "failed" - session_data["error"] = str( - data.get("message") or "Backend error message" - ) + session_data["error"] = str(data.get("message") or "Backend error message") break if session_data["status"] == "failed" and session_data["error"] is None: @@ -379,29 +342,18 @@ async def run_single_session( session_data["status"] = "failed" session_data["error"] = f"WebSocket connect/run failed: {exc}" - if ( - first_chunk_finish_epoch is not None - and last_chunk_finish_epoch is not None - and session_data["total_chunk_bytes"] > 0 - ): + if (first_chunk_finish_epoch is not None and last_chunk_finish_epoch is not None + and session_data["total_chunk_bytes"] > 0): duration_s = last_chunk_finish_epoch - first_chunk_finish_epoch if duration_s > 0: - session_data["session_goodput_mbps"] = ( - session_data["total_chunk_bytes"] * 8.0 / duration_s / 1_000_000.0 - ) + session_data["session_goodput_mbps"] = (session_data["total_chunk_bytes"] * 8.0 / duration_s / 1_000_000.0) - if ( - first_chunk_finish_epoch is not None - and first_media_segment_complete_epoch is not None - ): - session_data["first_chunk_before_first_media_complete"] = ( - first_chunk_finish_epoch < first_media_segment_complete_epoch - ) + if (first_chunk_finish_epoch is not None and first_media_segment_complete_epoch is not None): + session_data["first_chunk_before_first_media_complete"] = (first_chunk_finish_epoch + < first_media_segment_complete_epoch) session_data["close_ts_utc"] = utc_now_iso() - session_data["duration_ms"] = ( - time.monotonic() - session_start_monotonic - ) * 1000.0 + session_data["duration_ms"] = (time.monotonic() - session_start_monotonic) * 1000.0 return session_data @@ -412,14 +364,11 @@ async def run_worker_sessions( config: dict[str, Any], ) -> list[dict[str, Any]]: tasks = [ - asyncio.create_task( - run_single_session( - worker_id=worker_id, - worker_session_idx=idx, - config=config, - ) - ) - for idx in range(session_count) + asyncio.create_task(run_single_session( + worker_id=worker_id, + worker_session_idx=idx, + config=config, + )) for idx in range(session_count) ] if not tasks: return [] @@ -437,29 +386,23 @@ def worker_entry( try: ready_queue.put({"worker_id": worker_id, "status": "ready"}) start_event.wait() - sessions = asyncio.run( - run_worker_sessions( - worker_id=worker_id, - session_count=session_count, - config=config, - ) - ) - result_queue.put( - { - "worker_id": worker_id, - "status": "ok", - "sessions": sessions, - } - ) + sessions = asyncio.run(run_worker_sessions( + worker_id=worker_id, + session_count=session_count, + config=config, + )) + result_queue.put({ + "worker_id": worker_id, + "status": "ok", + "sessions": sessions, + }) except Exception as exc: - result_queue.put( - { - "worker_id": worker_id, - "status": "error", - "error": str(exc), - "traceback": traceback.format_exc(), - } - ) + result_queue.put({ + "worker_id": worker_id, + "status": "error", + "error": str(exc), + "traceback": traceback.format_exc(), + }) def build_summary( @@ -517,33 +460,22 @@ def build_summary( if len(all_chunk_finish_epochs) >= 2 and total_chunk_bytes > 0: duration_s = max(all_chunk_finish_epochs) - min(all_chunk_finish_epochs) if duration_s > 0: - global_goodput_mbps = ( - total_chunk_bytes * 8.0 / duration_s / 1_000_000.0 - ) + global_goodput_mbps = (total_chunk_bytes * 8.0 / duration_s / 1_000_000.0) - bucket_throughputs_mbps = [ - (bytes_count * 8.0) / 1_000_000.0 - for _, bytes_count in sorted(bucket_bytes.items()) - ] + bucket_throughputs_mbps = [(bytes_count * 8.0) / 1_000_000.0 for _, bytes_count in sorted(bucket_bytes.items())] bucket_stats = summarize_series(bucket_throughputs_mbps) - chunk_gap_threshold_breaches = [ - value for value in chunk_gaps if value >= chunk_gap_threshold_ms - ] + chunk_gap_threshold_breaches = [value for value in chunk_gaps if value >= chunk_gap_threshold_ms] non_success = len(sessions) - status_counts.get("success", 0) fail_reasons: list[str] = [] if non_success > 0: - fail_reasons.append( - f"{non_success} session(s) did not complete successfully." - ) + fail_reasons.append(f"{non_success} session(s) did not complete successfully.") if not chunk_gaps: fail_reasons.append("No chunk gap data collected.") if chunk_gap_threshold_breaches: - fail_reasons.append( - f"{len(chunk_gap_threshold_breaches)} chunk gap(s) were >= " - f"{chunk_gap_threshold_ms:.0f}ms." - ) + fail_reasons.append(f"{len(chunk_gap_threshold_breaches)} chunk gap(s) were >= " + f"{chunk_gap_threshold_ms:.0f}ms.") passed = len(fail_reasons) == 0 progressive_ratio = None @@ -554,20 +486,18 @@ def build_summary( "passed": passed, "fail_reasons": fail_reasons, "sessions": { - "total": len(sessions), - "success": status_counts.get("success", 0), - "failed": status_counts.get("failed", 0), - "timeout": status_counts.get("timeout", 0), - "protocol_error": status_counts.get("protocol_error", 0), - "other": ( - len(sessions) - - ( - status_counts.get("success", 0) - + status_counts.get("failed", 0) - + status_counts.get("timeout", 0) - + status_counts.get("protocol_error", 0) - ) - ), + "total": + len(sessions), + "success": + status_counts.get("success", 0), + "failed": + status_counts.get("failed", 0), + "timeout": + status_counts.get("timeout", 0), + "protocol_error": + status_counts.get("protocol_error", 0), + "other": (len(sessions) - (status_counts.get("success", 0) + status_counts.get("failed", 0) + + status_counts.get("timeout", 0) + status_counts.get("protocol_error", 0))), }, "chunk_gap_ms": { **chunk_gap_stats, @@ -606,51 +536,39 @@ def print_summary( bucket_bw = bandwidth["bucketed_1s"] print("=== LTX2 Realtime Stress Test Summary ===") - print( - "Run: " - f"url={run_info['url']} clients={run_info['clients']} " - f"processes={run_info['processes']} " - f"preset={run_info['preset_id']} " - f"curated_limit={run_info['curated_limit']}" - ) - print( - "Sessions: " - f"total={sessions['total']} success={sessions['success']} " - f"failed={sessions['failed']} timeout={sessions['timeout']} " - f"protocol_error={sessions['protocol_error']}" - ) - print( - "Chunk gap ms: " - f"min={format_num(chunk_gap['min'])} " - f"p50={format_num(chunk_gap['p50'])} " - f"p95={format_num(chunk_gap['p95'])} " - f"p99={format_num(chunk_gap['p99'])} " - f"max={format_num(chunk_gap['max'])} " - f"threshold={format_num(chunk_gap['threshold_ms'])} " - f"breaches={chunk_gap['breach_count']}" - ) - print( - "Queue wait ms: " - f"min={format_num(queue_wait['min'])} " - f"p50={format_num(queue_wait['p50'])} " - f"p95={format_num(queue_wait['p95'])} " - f"max={format_num(queue_wait['max'])}" - ) + print("Run: " + f"url={run_info['url']} clients={run_info['clients']} " + f"processes={run_info['processes']} " + f"preset={run_info['preset_id']} " + f"curated_limit={run_info['curated_limit']}") + print("Sessions: " + f"total={sessions['total']} success={sessions['success']} " + f"failed={sessions['failed']} timeout={sessions['timeout']} " + f"protocol_error={sessions['protocol_error']}") + print("Chunk gap ms: " + f"min={format_num(chunk_gap['min'])} " + f"p50={format_num(chunk_gap['p50'])} " + f"p95={format_num(chunk_gap['p95'])} " + f"p99={format_num(chunk_gap['p99'])} " + f"max={format_num(chunk_gap['max'])} " + f"threshold={format_num(chunk_gap['threshold_ms'])} " + f"breaches={chunk_gap['breach_count']}") + print("Queue wait ms: " + f"min={format_num(queue_wait['min'])} " + f"p50={format_num(queue_wait['p50'])} " + f"p95={format_num(queue_wait['p95'])} " + f"max={format_num(queue_wait['max'])}") ratio = progressive["ratio"] ratio_text = "n/a" if ratio is None else f"{ratio * 100:.2f}%" - print( - "Progressive streaming: " - f"{progressive['success_sessions']}/" - f"{progressive['eligible_sessions']} ({ratio_text})" - ) - print( - "Bandwidth Mbps: " - f"per_session_avg={format_num(per_session_bw['avg'])} " - f"per_session_p95={format_num(per_session_bw['p95'])} " - f"global={format_num(bandwidth['global_goodput_mbps'])} " - f"bucket_avg={format_num(bucket_bw['avg_mbps'])} " - f"bucket_peak={format_num(bucket_bw['peak_mbps'])}" - ) + print("Progressive streaming: " + f"{progressive['success_sessions']}/" + f"{progressive['eligible_sessions']} ({ratio_text})") + print("Bandwidth Mbps: " + f"per_session_avg={format_num(per_session_bw['avg'])} " + f"per_session_p95={format_num(per_session_bw['p95'])} " + f"global={format_num(bandwidth['global_goodput_mbps'])} " + f"bucket_avg={format_num(bucket_bw['avg_mbps'])} " + f"bucket_peak={format_num(bucket_bw['peak_mbps'])}") print(f"VERDICT: {'PASS' if summary['passed'] else 'FAIL'}") if summary["fail_reasons"]: print("Fail reasons:") @@ -670,10 +588,8 @@ def distribute_sessions(total_clients: int, process_count: int) -> list[int]: def run_stress(args: argparse.Namespace) -> tuple[dict[str, Any], int]: if websockets is None: - raise RuntimeError( - "Missing dependency: websockets. Install it before running this " - "stress test." - ) + raise RuntimeError("Missing dependency: websockets. Install it before running this " + "stress test.") preset_file = Path(args.preset_file).expanduser().resolve() selected_preset_id, curated_prompts, total_prompt_count = load_curated_prompts( @@ -735,13 +651,8 @@ def run_stress(args: argparse.Namespace) -> tuple[dict[str, Any], int]: start_event.set() - result_deadline = ( - time.monotonic() - + args.connect_timeout_s - + args.session_timeout_s - + args.post_complete_wait_s - + 180.0 - ) + result_deadline = (time.monotonic() + args.connect_timeout_s + args.session_timeout_s + + args.post_complete_wait_s + 180.0) worker_results: list[dict[str, Any]] = [] while len(worker_results) < len(processes): timeout_s = max(0.1, result_deadline - time.monotonic()) @@ -765,24 +676,20 @@ def run_stress(args: argparse.Namespace) -> tuple[dict[str, Any], int]: if result.get("status") == "ok": sessions.extend(result.get("sessions", [])) else: - worker_errors.append( - { - "worker_id": result.get("worker_id"), - "error": result.get("error"), - "traceback": result.get("traceback"), - } - ) + worker_errors.append({ + "worker_id": result.get("worker_id"), + "error": result.get("error"), + "traceback": result.get("traceback"), + }) received_workers = {result.get("worker_id") for result in worker_results} expected_workers = set(range(len(processes))) missing_workers = sorted(expected_workers - received_workers) for worker_id in missing_workers: - worker_errors.append( - { - "worker_id": worker_id, - "error": "No worker result received.", - } - ) + worker_errors.append({ + "worker_id": worker_id, + "error": "No worker result received.", + }) run_end_epoch = time.time() run_end_iso = iso_from_epoch(run_end_epoch) @@ -795,9 +702,8 @@ def run_stress(args: argparse.Namespace) -> tuple[dict[str, Any], int]: if worker_errors: summary["passed"] = False - summary["fail_reasons"] = list(summary["fail_reasons"]) + [ - f"{len(worker_errors)} worker error(s) occurred." - ] + summary["fail_reasons"] = list( + summary["fail_reasons"]) + [f"{len(worker_errors)} worker error(s) occurred."] output_payload = { "run_info": { @@ -833,9 +739,7 @@ def run_stress(args: argparse.Namespace) -> tuple[dict[str, Any], int]: def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Multiprocess realtime stress test for LTX2 streaming.", - ) + parser = argparse.ArgumentParser(description="Multiprocess realtime stress test for LTX2 streaming.", ) parser.add_argument( "-u", "--url", diff --git a/apps/dreamverse/dreamverse/tests/test_session_init_image.py b/apps/dreamverse/dreamverse/tests/test_session_init_image.py index 52e351c0c..26b16db41 100644 --- a/apps/dreamverse/dreamverse/tests/test_session_init_image.py +++ b/apps/dreamverse/dreamverse/tests/test_session_init_image.py @@ -47,13 +47,11 @@ def test_persist_session_init_image_returns_none_when_missing_data(): def test_persist_session_init_image_rejects_unsupported_mime(): with pytest.raises(ValueError, match="PNG, JPEG, or WebP"): - persist_session_init_image( - { - "name": "frame.gif", - "mime_type": "image/gif", - "data_url": "data:image/gif;base64,R0lGODlhAQABAAAAACw=", - } - ) + persist_session_init_image({ + "name": "frame.gif", + "mime_type": "image/gif", + "data_url": "data:image/gif;base64,R0lGODlhAQABAAAAACw=", + }) def test_persist_session_init_image_rejects_large_payload(monkeypatch): @@ -66,10 +64,8 @@ def test_persist_session_init_image_rejects_large_payload(monkeypatch): monkeypatch.setattr(base64, "b64decode", fake_b64decode) with pytest.raises(ValueError, match="15 MB or smaller"): - persist_session_init_image( - { - "name": "frame.png", - "mime_type": "image/png", - "data_url": data_url, - } - ) + persist_session_init_image({ + "name": "frame.png", + "mime_type": "image/png", + "data_url": data_url, + }) diff --git a/apps/dreamverse/dreamverse/tests/test_session_logging.py b/apps/dreamverse/dreamverse/tests/test_session_logging.py index 726f70e4b..ead39b924 100644 --- a/apps/dreamverse/dreamverse/tests/test_session_logging.py +++ b/apps/dreamverse/dreamverse/tests/test_session_logging.py @@ -14,7 +14,6 @@ from fastapi import WebSocketDisconnect from PIL import Image import pytest - os.environ.setdefault("CEREBRAS_API_KEY", "dummy") os.environ.setdefault("GROQ_API_KEY", "dummy") @@ -30,6 +29,7 @@ pytestmark = pytest.mark.gpu class _FakeSessionLogger: + def __init__(self): self.events: list[dict[str, object]] = [] @@ -50,6 +50,7 @@ class _FakeSessionLogger: class _FakeWebSocket: + def __init__(self, messages: list[tuple[float, dict[str, object]]]): self._messages = messages self._index = 0 @@ -79,6 +80,7 @@ class _FakeWebSocket: class _FakeSlot: + def __init__(self, step_delay_s: float = 0.0): self.shared_stream_buffer = None self._stream_queues: dict[str, asyncio.Queue] = {} @@ -112,32 +114,36 @@ class _FakeSlot: "reset_conditioning": reset_conditioning, }) queue = self._stream_queues[client_id] - await queue.put(MediaInit( - user_id=client_id, - segment_idx=segment_idx, - stream_id=f"stream-{segment_idx}", - mime='video/mp4; codecs="avc1.640028,mp4a.40.2"', - uses_shared_buffer=False, - )) - await queue.put(MediaChunk( - user_id=client_id, - segment_idx=segment_idx, - stream_id=f"stream-{segment_idx}", - chunk=b"abc123", - uses_shared_buffer=False, - )) - await queue.put(MediaComplete( - user_id=client_id, - segment_idx=segment_idx, - stream_id=f"stream-{segment_idx}", - chunks=1, - )) + await queue.put( + MediaInit( + user_id=client_id, + segment_idx=segment_idx, + stream_id=f"stream-{segment_idx}", + mime='video/mp4; codecs="avc1.640028,mp4a.40.2"', + uses_shared_buffer=False, + )) + await queue.put( + MediaChunk( + user_id=client_id, + segment_idx=segment_idx, + stream_id=f"stream-{segment_idx}", + chunk=b"abc123", + uses_shared_buffer=False, + )) + await queue.put( + MediaComplete( + user_id=client_id, + segment_idx=segment_idx, + stream_id=f"stream-{segment_idx}", + chunks=1, + )) if self._step_delay_s > 0: await asyncio.sleep(self._step_delay_s) return {"e2e_latency_ms": 10.0} class _FakeGPUPool: + def __init__(self, step_delay_s: float = 0.0): self._slot = _FakeSlot(step_delay_s=step_delay_s) @@ -165,6 +171,7 @@ def _make_png_data_url() -> str: class _FakePromptEnhancer: + def set_provider(self, provider: object, preferred_model: object = None): del provider, preferred_model return { @@ -246,14 +253,13 @@ class _FakePromptEnhancer: latency_ms=23.0, rollout_id="rewrite_rollout", rollout_label="Rewrite Rollout", - raw_response_text=( - '{"id":"rewrite_rollout","label":"Rewrite Rollout",' - '"segment_prompts":["rewrite a","rewrite b"]}' - ), + raw_response_text=('{"id":"rewrite_rollout","label":"Rewrite Rollout",' + '"segment_prompts":["rewrite a","rewrite b"]}'), ) class _FallbackPromptEnhancer(_FakePromptEnhancer): + async def enhance_prompt( self, raw_prompt: str, @@ -285,6 +291,7 @@ class _FallbackPromptEnhancer(_FakePromptEnhancer): class _UnsafeEnhancedPromptEnhancer(_FakePromptEnhancer): + async def enhance_prompt( self, raw_prompt: str, @@ -316,11 +323,11 @@ class _UnsafeEnhancedPromptEnhancer(_FakePromptEnhancer): class _FakePromptSafetyFilter: + def __init__(self, blocked_prompts: dict[str, str]): self._blocked_prompts = { str(prompt).strip(): str(error) - for prompt, error in blocked_prompts.items() - if str(prompt).strip() + for prompt, error in blocked_prompts.items() if str(prompt).strip() } def get_prompt_safety_error(self, prompt: str) -> str | None: @@ -334,9 +341,7 @@ class _FakePromptSafetyFilter: return None -def test_session_event_logger_initializes_hostname_folder_and_utc_filename( - tmp_path, -): +def test_session_event_logger_initializes_hostname_folder_and_utc_filename(tmp_path, ): logger = SessionEventLogger(tmp_path) assert logger.directory == tmp_path / socket.gethostname() assert logger.directory.is_dir() @@ -376,29 +381,32 @@ def test_websocket_flow_emits_required_session_events(): runtime.prompt_enhancer = _FakePromptEnhancer() runtime.session_event_logger = fake_logger - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "curated_prompts": ["segment one prompt"], - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - (0.01, {"type": "append_prompt", "prompt": "make it rainy"}), - ( - 0.01, - { - "type": "rewrite_seed_prompts", - "rewrite_instruction": "more dramatic", - }, - ), - (0.30, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "curated_prompts": ["segment one prompt"], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.01, { + "type": "append_prompt", + "prompt": "make it rainy" + }), + ( + 0.01, + { + "type": "rewrite_seed_prompts", + "rewrite_instruction": "more dramatic", + }, + ), + (0.30, { + "type": "leave" + }), + ]) asyncio.run(server_main.websocket_endpoint(ws)) @@ -409,20 +417,12 @@ def test_websocket_flow_emits_required_session_events(): assert "segment_start" in event_names assert "segment_complete" in event_names - rewrite_events = [ - event - for event in fake_logger.events - if event["event"] == "rewrite_done" - ] + rewrite_events = [event for event in fake_logger.events if event["event"] == "rewrite_done"] rewrite_kinds = {event.get("kind") for event in rewrite_events} assert "seed_rewrite" in rewrite_kinds assert "enhance_prompt" in rewrite_kinds - segment_complete = next( - event - for event in fake_logger.events - if event["event"] == "segment_complete" - ) + segment_complete = next(event for event in fake_logger.events if event["event"] == "segment_complete") latency = segment_complete.get("latency_ms") assert isinstance(latency, dict) assert set(latency.keys()) == { @@ -452,38 +452,38 @@ def test_simple_generate_reuses_session_and_resets_conditioning(): runtime.session_event_logger = fake_logger runtime.prompt_safety_filter = None - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "simple_prompt_1", - "curated_prompts": ["a selected prompt"], - "single_clip_mode": True, - "enhancement_enabled": False, - "auto_extension_enabled": False, - "loop_generation_enabled": False, + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "simple_prompt_1", + "curated_prompts": ["a selected prompt"], + "single_clip_mode": True, + "enhancement_enabled": False, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "simple_generate", + "preset_id": "simple_custom_prompt", + "prompt_id": "simple_custom_prompt", + "prompt": "a fresh custom prompt", + "enhancement_enabled": True, + "initial_image": { + "name": "frame.png", + "mime_type": "image/png", + "data_url": _make_png_data_url(), }, - ), - ( - 0.01, - { - "type": "simple_generate", - "preset_id": "simple_custom_prompt", - "prompt_id": "simple_custom_prompt", - "prompt": "a fresh custom prompt", - "enhancement_enabled": True, - "initial_image": { - "name": "frame.png", - "mime_type": "image/png", - "data_url": _make_png_data_url(), - }, - }, - ), - (0.50, {"type": "leave"}), - ] - ) + }, + ), + (0.50, { + "type": "leave" + }), + ]) asyncio.run(server_main.websocket_endpoint(ws)) @@ -498,38 +498,17 @@ def test_simple_generate_reuses_session_and_resets_conditioning(): event_names = [event["event"] for event in fake_logger.events] assert "simple_generate" in event_names - stream_start_events = [ - payload - for payload in ws.sent_json - if payload.get("type") == "ltx2_stream_start" - ] + stream_start_events = [payload for payload in ws.sent_json if payload.get("type") == "ltx2_stream_start"] assert len(stream_start_events) == 2 - assert all( - payload.get("generation_segment_cap") == 0 - for payload in stream_start_events - ) - stream_complete_events = [ - payload - for payload in ws.sent_json - if payload.get("type") == "ltx2_stream_complete" - ] + assert all(payload.get("generation_segment_cap") == 0 for payload in stream_start_events) + stream_complete_events = [payload for payload in ws.sent_json if payload.get("type") == "ltx2_stream_complete"] assert len(stream_complete_events) == 2 - assert "prompt_sources_blocked" not in [ - payload.get("type") for payload in ws.sent_json - ] + assert "prompt_sources_blocked" not in [payload.get("type") for payload in ws.sent_json] - segment_start_events = [ - payload - for payload in ws.sent_json - if payload.get("type") == "ltx2_segment_start" - ] + segment_start_events = [payload for payload in ws.sent_json if payload.get("type") == "ltx2_segment_start"] assert [payload["segment_idx"] for payload in segment_start_events] == [1, 1] - step_complete_events = [ - payload - for payload in ws.sent_json - if payload.get("type") == "step_complete" - ] + step_complete_events = [payload for payload in ws.sent_json if payload.get("type") == "step_complete"] assert len(step_complete_events) == 2 assert all( set(payload.get("latency_ms", {}).keys()) == { @@ -537,9 +516,7 @@ def test_simple_generate_reuses_session_and_resets_conditioning(): "worker_e2e", "main_user_step", "overhead", - } - for payload in step_complete_events - ) + } for payload in step_complete_events) finally: runtime.gpu_pool = old_gpu_pool runtime.prompt_enhancer = old_prompt_enhancer @@ -562,46 +539,42 @@ def test_single5s_ignores_global_segment_cap(): runtime.prompt_safety_filter = None session_controller.GENERATION_SEGMENT_CAP = 1 - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "simple_prompt_1", - "curated_prompts": ["a selected prompt"], - "single_clip_mode": True, - "enhancement_enabled": False, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - ( - 0.01, - { - "type": "simple_generate", - "preset_id": "simple_custom_prompt", - "prompt_id": "simple_custom_prompt", - "prompt": "a fresh custom prompt", - "enhancement_enabled": True, - "initial_image": None, - }, - ), - (0.50, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "simple_prompt_1", + "curated_prompts": ["a selected prompt"], + "single_clip_mode": True, + "enhancement_enabled": False, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "simple_generate", + "preset_id": "simple_custom_prompt", + "prompt_id": "simple_custom_prompt", + "prompt": "a fresh custom prompt", + "enhancement_enabled": True, + "initial_image": None, + }, + ), + (0.50, { + "type": "leave" + }), + ]) asyncio.run(server_main.websocket_endpoint(ws)) assert [call["segment_idx"] for call in fake_pool._slot.calls] == [1, 1] assert all( - payload.get("generation_segment_cap") == 0 - for payload in ws.sent_json - if payload.get("type") == "ltx2_stream_start" - ) - assert "generation_cap_reached" not in [ - payload.get("type") for payload in ws.sent_json - ] + payload.get("generation_segment_cap") == 0 for payload in ws.sent_json + if payload.get("type") == "ltx2_stream_start") + assert "generation_cap_reached" not in [payload.get("type") for payload in ws.sent_json] finally: runtime.gpu_pool = old_gpu_pool runtime.prompt_enhancer = old_prompt_enhancer @@ -625,29 +598,29 @@ def test_regular_segment_cap_completes_stream_and_allows_rewrite_rollout_restart runtime.prompt_safety_filter = None session_controller.GENERATION_SEGMENT_CAP = 1 - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "test_preset", - "curated_prompts": ["segment one prompt"], - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - ( - 0.20, - { - "type": "rewrite_seed_prompts", - "rewrite_instruction": "start a new rollout", - }, - ), - (0.50, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "test_preset", + "curated_prompts": ["segment one prompt"], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.20, + { + "type": "rewrite_seed_prompts", + "rewrite_instruction": "start a new rollout", + }, + ), + (0.50, { + "type": "leave" + }), + ]) asyncio.run(server_main.websocket_endpoint(ws)) @@ -682,32 +655,32 @@ def test_rewrite_during_active_segment_restarts_from_rewritten_seed_window(): runtime.session_event_logger = fake_logger runtime.prompt_safety_filter = None - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "test_preset", - "curated_prompts": [ - "segment one prompt", - "segment two prompt", - ], - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - ( - 0.01, - { - "type": "rewrite_seed_prompts", - "rewrite_instruction": "restart from rewrite", - }, - ), - (0.45, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "test_preset", + "curated_prompts": [ + "segment one prompt", + "segment two prompt", + ], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "rewrite_seed_prompts", + "rewrite_instruction": "restart from rewrite", + }, + ), + (0.45, { + "type": "leave" + }), + ]) asyncio.run(server_main.websocket_endpoint(ws)) @@ -715,23 +688,13 @@ def test_rewrite_during_active_segment_restarts_from_rewritten_seed_window(): "segment one prompt", "rewrite a", ] - assert all( - call["prompt"] != "segment two prompt" - for call in fake_pool._slot.calls[1:] - ) + assert all(call["prompt"] != "segment two prompt" for call in fake_pool._slot.calls[1:]) assert fake_pool._slot.calls[1]["segment_idx"] == 1 assert fake_pool._slot.calls[1]["reset_conditioning"] is True - reset_events = [ - payload - for payload in ws.sent_json - if payload.get("type") == "seed_prompts_reset_applied" - ] + reset_events = [payload for payload in ws.sent_json if payload.get("type") == "seed_prompts_reset_applied"] assert reset_events - assert any( - payload.get("reason") == "rewrite_during_generation" - for payload in reset_events - ) + assert any(payload.get("reason") == "rewrite_during_generation" for payload in reset_events) finally: runtime.gpu_pool = old_gpu_pool runtime.prompt_enhancer = old_prompt_enhancer @@ -740,7 +703,9 @@ def test_rewrite_during_active_segment_restarts_from_rewritten_seed_window(): def test_initial_custom_rollout_prompt_generates_seed_window_before_streaming(): + class _InitialRolloutPromptEnhancer(_FakePromptEnhancer): + def __init__(self): self.rewrite_calls: list[dict[str, object]] = [] @@ -783,11 +748,9 @@ def test_initial_custom_rollout_prompt_generates_seed_window_before_streaming(): latency_ms=23.0, rollout_id=preset_id or "custom_editable", rollout_label=preset_label or "Custom rollout", - raw_response_text=( - '{"id":"custom_editable","label":"Custom rollout",' - '"segment_prompts":["rewrite 1","rewrite 2","rewrite 3",' - '"rewrite 4","rewrite 5","rewrite 6"]}' - ), + raw_response_text=('{"id":"custom_editable","label":"Custom rollout",' + '"segment_prompts":["rewrite 1","rewrite 2","rewrite 3",' + '"rewrite 4","rewrite 5","rewrite 6"]}'), ) old_gpu_pool = runtime.gpu_pool @@ -803,27 +766,27 @@ def test_initial_custom_rollout_prompt_generates_seed_window_before_streaming(): runtime.session_event_logger = fake_logger runtime.prompt_safety_filter = None - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "custom_editable", - "preset_label": "Custom rollout", - "curated_prompts": [], - "initial_rollout_prompt": "A moonbase corridor thriller with flooding", - "rewrite_model": "gpt-4.1-mini", - "rewrite_temperature": 0.4, - "rewrite_window_system_prompt": "Session-specific rewrite prompt", - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - (0.20, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "custom_editable", + "preset_label": "Custom rollout", + "curated_prompts": [], + "initial_rollout_prompt": "A moonbase corridor thriller with flooding", + "rewrite_model": "gpt-4.1-mini", + "rewrite_temperature": 0.4, + "rewrite_window_system_prompt": "Session-specific rewrite prompt", + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.20, { + "type": "leave" + }), + ]) asyncio.run(server_main.websocket_endpoint(ws)) @@ -856,7 +819,9 @@ def test_initial_custom_rollout_prompt_generates_seed_window_before_streaming(): def test_initial_custom_rollout_prompt_uses_dedicated_user_prompt_by_default(): + class _InitialRolloutPromptEnhancer(_FakePromptEnhancer): + def __init__(self): self.rewrite_calls: list[dict[str, object]] = [] @@ -898,11 +863,9 @@ def test_initial_custom_rollout_prompt_uses_dedicated_user_prompt_by_default(): latency_ms=23.0, rollout_id=preset_id or "custom_editable", rollout_label=preset_label or "Custom rollout", - raw_response_text=( - '{"id":"custom_editable","label":"Custom rollout",' - '"segment_prompts":["rewrite 1","rewrite 2","rewrite 3",' - '"rewrite 4","rewrite 5","rewrite 6"]}' - ), + raw_response_text=('{"id":"custom_editable","label":"Custom rollout",' + '"segment_prompts":["rewrite 1","rewrite 2","rewrite 3",' + '"rewrite 4","rewrite 5","rewrite 6"]}'), ) old_gpu_pool = runtime.gpu_pool @@ -918,35 +881,32 @@ def test_initial_custom_rollout_prompt_uses_dedicated_user_prompt_by_default(): runtime.session_event_logger = fake_logger runtime.prompt_safety_filter = None - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "custom_editable", - "preset_label": "Custom rollout", - "curated_prompts": [], - "initial_rollout_prompt": "A moonbase corridor thriller with flooding", - "rewrite_model": "gpt-4.1-mini", - "rewrite_temperature": 0.4, - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - (0.20, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "custom_editable", + "preset_label": "Custom rollout", + "curated_prompts": [], + "initial_rollout_prompt": "A moonbase corridor thriller with flooding", + "rewrite_model": "gpt-4.1-mini", + "rewrite_temperature": 0.4, + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.20, { + "type": "leave" + }), + ]) asyncio.run(server_main.websocket_endpoint(ws)) assert len(fake_enhancer.rewrite_calls) == 1 assert fake_enhancer.rewrite_calls[0]["snapshot_prompts"] == [] - assert ( - fake_enhancer.rewrite_calls[0]["system_prompt_override"] - == "new rollout rewrite system prompt" - ) + assert (fake_enhancer.rewrite_calls[0]["system_prompt_override"] == "new rollout rewrite system prompt") finally: runtime.gpu_pool = old_gpu_pool runtime.prompt_enhancer = old_prompt_enhancer @@ -955,7 +915,9 @@ def test_initial_custom_rollout_prompt_uses_dedicated_user_prompt_by_default(): def test_initial_custom_rollout_prompt_prefers_session_specific_user_prompt(): + class _InitialRolloutPromptEnhancer(_FakePromptEnhancer): + def __init__(self): self.rewrite_calls: list[dict[str, object]] = [] @@ -997,11 +959,9 @@ def test_initial_custom_rollout_prompt_prefers_session_specific_user_prompt(): latency_ms=23.0, rollout_id=preset_id or "custom_editable", rollout_label=preset_label or "Custom rollout", - raw_response_text=( - '{"id":"custom_editable","label":"Custom rollout",' - '"segment_prompts":["rewrite 1","rewrite 2","rewrite 3",' - '"rewrite 4","rewrite 5","rewrite 6"]}' - ), + raw_response_text=('{"id":"custom_editable","label":"Custom rollout",' + '"segment_prompts":["rewrite 1","rewrite 2","rewrite 3",' + '"rewrite 4","rewrite 5","rewrite 6"]}'), ) old_gpu_pool = runtime.gpu_pool @@ -1017,28 +977,28 @@ def test_initial_custom_rollout_prompt_prefers_session_specific_user_prompt(): runtime.session_event_logger = fake_logger runtime.prompt_safety_filter = None - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "custom_editable", - "preset_label": "Custom rollout", - "curated_prompts": [], - "initial_rollout_prompt": "A moonbase corridor thriller with flooding", - "rewrite_model": "gpt-4.1-mini", - "rewrite_temperature": 0.4, - "rewrite_window_system_prompt": "Window rewrite prompt", - "rewrite_user_system_prompt": "User rewrite prompt", - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - (0.20, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "custom_editable", + "preset_label": "Custom rollout", + "curated_prompts": [], + "initial_rollout_prompt": "A moonbase corridor thriller with flooding", + "rewrite_model": "gpt-4.1-mini", + "rewrite_temperature": 0.4, + "rewrite_window_system_prompt": "Window rewrite prompt", + "rewrite_user_system_prompt": "User rewrite prompt", + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.20, { + "type": "leave" + }), + ]) asyncio.run(server_main.websocket_endpoint(ws)) @@ -1067,51 +1027,40 @@ def test_single5s_raw_prompt_block_prevents_generation(): "blocked raw prompt": "Test raw prompt blocked.", }) - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "simple_custom_prompt", - "curated_prompts": [], - "single_clip_mode": True, - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - ( - 0.01, - { - "type": "append_prompt", - "prompt_id": "simple_custom_prompt", - "prompt": "blocked raw prompt", - }, - ), - (0.20, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "simple_custom_prompt", + "curated_prompts": [], + "single_clip_mode": True, + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "append_prompt", + "prompt_id": "simple_custom_prompt", + "prompt": "blocked raw prompt", + }, + ), + (0.20, { + "type": "leave" + }), + ]) asyncio.run(server_main.websocket_endpoint(ws)) assert fake_pool._slot.calls == [] - assert { - payload["type"] for payload in ws.sent_json - }.isdisjoint({"ltx2_segment_start", "ltx2_segment_complete"}) - error_payloads = [ - payload for payload in ws.sent_json if payload.get("type") == "error" - ] - assert any( - payload.get("message") == "Test raw prompt blocked." - for payload in error_payloads - ) + assert {payload["type"] for payload in ws.sent_json}.isdisjoint({"ltx2_segment_start", "ltx2_segment_complete"}) + error_payloads = [payload for payload in ws.sent_json if payload.get("type") == "error"] + assert any(payload.get("message") == "Test raw prompt blocked." for payload in error_payloads) - prompt_blocked_events = [ - event - for event in fake_logger.events - if event["event"] == "prompt_blocked" - ] + prompt_blocked_events = [event for event in fake_logger.events if event["event"] == "prompt_blocked"] assert len(prompt_blocked_events) == 1 assert prompt_blocked_events[0]["kind"] == "user_raw" assert prompt_blocked_events[0]["raw_prompt"] == "blocked raw prompt" @@ -1136,62 +1085,48 @@ def test_single5s_enhanced_prompt_block_prevents_generation(): runtime.prompt_enhancer = _UnsafeEnhancedPromptEnhancer() runtime.session_event_logger = fake_logger runtime.prompt_safety_filter = _FakePromptSafetyFilter({ - "unsafe enhanced prompt": "Test enhanced prompt blocked.", + "unsafe enhanced prompt": + "Test enhanced prompt blocked.", }) - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "simple_custom_prompt", - "curated_prompts": [], - "single_clip_mode": True, - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - ( - 0.01, - { - "type": "append_prompt", - "prompt_id": "simple_custom_prompt", - "prompt": "safe raw prompt", - }, - ), - (0.20, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "simple_custom_prompt", + "curated_prompts": [], + "single_clip_mode": True, + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "append_prompt", + "prompt_id": "simple_custom_prompt", + "prompt": "safe raw prompt", + }, + ), + (0.20, { + "type": "leave" + }), + ]) asyncio.run(server_main.websocket_endpoint(ws)) assert fake_pool._slot.calls == [] - assert { - payload["type"] for payload in ws.sent_json - }.isdisjoint({"ltx2_segment_start", "ltx2_segment_complete"}) - error_payloads = [ - payload for payload in ws.sent_json if payload.get("type") == "error" - ] - assert any( - payload.get("message") == "Test enhanced prompt blocked." - for payload in error_payloads - ) + assert {payload["type"] for payload in ws.sent_json}.isdisjoint({"ltx2_segment_start", "ltx2_segment_complete"}) + error_payloads = [payload for payload in ws.sent_json if payload.get("type") == "error"] + assert any(payload.get("message") == "Test enhanced prompt blocked." for payload in error_payloads) - rewrite_events = [ - event - for event in fake_logger.events - if event["event"] == "rewrite_done" - ] + rewrite_events = [event for event in fake_logger.events if event["event"] == "rewrite_done"] assert len(rewrite_events) == 1 assert rewrite_events[0]["kind"] == "enhance_prompt" - prompt_blocked_events = [ - event - for event in fake_logger.events - if event["event"] == "prompt_blocked" - ] + prompt_blocked_events = [event for event in fake_logger.events if event["event"] == "prompt_blocked"] assert len(prompt_blocked_events) == 1 assert prompt_blocked_events[0]["kind"] == "user_enhanced" assert prompt_blocked_events[0]["output_prompt"] == "unsafe enhanced prompt" @@ -1216,57 +1151,44 @@ def test_single5s_enhancement_fallback_does_not_start_generation(): runtime.session_event_logger = fake_logger runtime.prompt_safety_filter = None - ws = _FakeWebSocket( - [ - ( - 0.0, - { - "type": "session_init_v2", - "preset_id": "simple_custom_prompt", - "curated_prompts": [], - "single_clip_mode": True, - "enhancement_enabled": True, - "auto_extension_enabled": False, - "loop_generation_enabled": False, - }, - ), - ( - 0.01, - { - "type": "append_prompt", - "prompt_id": "simple_custom_prompt", - "prompt": "prompt that times out", - }, - ), - (0.20, {"type": "leave"}), - ] - ) + ws = _FakeWebSocket([ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "simple_custom_prompt", + "curated_prompts": [], + "single_clip_mode": True, + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "append_prompt", + "prompt_id": "simple_custom_prompt", + "prompt": "prompt that times out", + }, + ), + (0.20, { + "type": "leave" + }), + ]) asyncio.run(server_main.websocket_endpoint(ws)) assert fake_pool._slot.calls == [] - assert { - payload["type"] for payload in ws.sent_json - }.isdisjoint({"ltx2_segment_start", "ltx2_segment_complete"}) + assert {payload["type"] for payload in ws.sent_json}.isdisjoint({"ltx2_segment_start", "ltx2_segment_complete"}) - fallback_payloads = [ - payload - for payload in ws.sent_json - if payload.get("type") == "prompt_fallback_used" - ] + fallback_payloads = [payload for payload in ws.sent_json if payload.get("type") == "prompt_fallback_used"] assert len(fallback_payloads) == 1 assert fallback_payloads[0]["prompt_id"] == "simple_custom_prompt" - assert ( - fallback_payloads[0]["error"] - == PROMPT_EXTENSION_FAILURE_USER_MESSAGE - ) + assert (fallback_payloads[0]["error"] == PROMPT_EXTENSION_FAILURE_USER_MESSAGE) assert fallback_payloads[0]["source"] == "user_enhancement_failed" - rewrite_events = [ - event - for event in fake_logger.events - if event["event"] == "rewrite_done" - ] + rewrite_events = [event for event in fake_logger.events if event["event"] == "rewrite_done"] assert len(rewrite_events) == 1 assert rewrite_events[0]["kind"] == "enhance_prompt" assert all(event["event"] != "segment_start" for event in fake_logger.events) diff --git a/apps/dreamverse/scripts/modal/modal_app.py b/apps/dreamverse/scripts/modal/modal_app.py index b0919a74a..c710caa28 100644 --- a/apps/dreamverse/scripts/modal/modal_app.py +++ b/apps/dreamverse/scripts/modal/modal_app.py @@ -8,12 +8,10 @@ import modal IMAGE = os.environ.get("DREAMVERSE_IMAGE") if not IMAGE: - raise RuntimeError( - "DREAMVERSE_IMAGE is required. Set it to a published SHA-specific Dreamverse image, " - "for example a dreamverse-backend-cuda13.0.0-sha-* tag or a " - "dreamverse-ui-cuda13.0.0-sha-* tag if serving the static UI. " - "CUDA 12 / cu126 images use the corresponding cuda12.6.3 tag." - ) + raise RuntimeError("DREAMVERSE_IMAGE is required. Set it to a published SHA-specific Dreamverse image, " + "for example a dreamverse-backend-cuda13.0.0-sha-* tag or a " + "dreamverse-ui-cuda13.0.0-sha-* tag if serving the static UI. " + "CUDA 12 / cu126 images use the corresponding cuda12.6.3 tag.") # ``@modal.web_server`` invokes ``serve()`` directly and bypasses the image # ENTRYPOINT (``docker/docker_entrypoint.sh``). That entrypoint normally @@ -65,14 +63,10 @@ def serve(): # ``or ""`` collapses ``None`` (unset) into an empty string, ``.strip()`` # collapses whitespace-only values (e.g. ``" "``) — both should be # treated as missing. - missing = [ - k for k in _REQUIRED_SECRET_KEYS - if not (os.environ.get(k) or "").strip() - ] + missing = [k for k in _REQUIRED_SECRET_KEYS if not (os.environ.get(k) or "").strip()] if missing: - raise RuntimeError( - "dreamverse-api-keys secret is missing required entries: " - f"{', '.join(missing)}. Add them with `modal secret create " - "dreamverse-api-keys ... --force` and redeploy " - "(see apps/dreamverse/scripts/modal/README.md).") + raise RuntimeError("dreamverse-api-keys secret is missing required entries: " + f"{', '.join(missing)}. Add them with `modal secret create " + "dreamverse-api-keys ... --force` and redeploy " + "(see apps/dreamverse/scripts/modal/README.md).") subprocess.Popen(["dreamverse-server", "--host", "0.0.0.0", "--port", "8009"]) diff --git a/apps/fastvideo_studio/tests/test_gpu.py b/apps/fastvideo_studio/tests/test_gpu.py index 63d919390..3f670b01a 100644 --- a/apps/fastvideo_studio/tests/test_gpu.py +++ b/apps/fastvideo_studio/tests/test_gpu.py @@ -74,10 +74,8 @@ def test_snapshot_shapes_devices(monkeypatch: pytest.MonkeyPatch) -> None: } -def test_snapshot_tolerates_missing_sensors( - monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setitem(sys.modules, "pynvml", - _make_fake_pynvml(broken_sensors=True)) +def test_snapshot_tolerates_missing_sensors(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(sys.modules, "pynvml", _make_fake_pynvml(broken_sensors=True)) snap = gpu_mod.get_gpu_snapshot() assert snap["available"] is True g = snap["gpus"][0] @@ -86,8 +84,7 @@ def test_snapshot_tolerates_missing_sensors( assert g["power_limit_watts"] is None -def test_snapshot_reports_nvml_failure( - monkeypatch: pytest.MonkeyPatch) -> None: +def test_snapshot_reports_nvml_failure(monkeypatch: pytest.MonkeyPatch) -> None: fake = _make_fake_pynvml() fake.nvmlInit = lambda: (_ for _ in ()).throw(_NVMLError("driver gone")) monkeypatch.setitem(sys.modules, "pynvml", fake) diff --git a/apps/fastvideo_studio/tests/test_training_config.py b/apps/fastvideo_studio/tests/test_training_config.py index db595e743..ffa2d8fd4 100644 --- a/apps/fastvideo_studio/tests/test_training_config.py +++ b/apps/fastvideo_studio/tests/test_training_config.py @@ -130,8 +130,7 @@ def test_dmd_builds_three_role_models_and_method_knobs() -> None: def test_dmd_vsa_maps_to_training_vsa_sparsity() -> None: - config = build_training_config( - _job("dmd_t2v", dmd_use_vsa=True, dmd_vsa_sparsity=0.9), "out") + config = build_training_config(_job("dmd_t2v", dmd_use_vsa=True, dmd_vsa_sparsity=0.9), "out") assert config["training"]["vsa"]["sparsity"] == 0.9 @@ -161,31 +160,27 @@ def test_validation_callback_only_when_file_given() -> None: without = build_training_config(_job("full_t2v"), "out") assert "validation" not in without["callbacks"] - with_file = build_training_config( - _job("full_t2v", validation_dataset_file="val.json"), "out") + with_file = build_training_config(_job("full_t2v", validation_dataset_file="val.json"), "out") validation = with_file["callbacks"]["validation"] assert validation["dataset_file"] == "val.json" assert validation["pipeline_target"].endswith(".WanPipeline") assert validation["sampling_steps"] == [50] - dmd = build_training_config( - _job("dmd_t2v", validation_dataset_file="val.json"), "out") + dmd = build_training_config(_job("dmd_t2v", validation_dataset_file="val.json"), "out") validation = dmd["callbacks"]["validation"] assert validation["pipeline_target"].endswith(".WanDMDPipeline") assert validation["sampling_steps"] == [3] assert validation["sampling_timesteps"] == [1000, 757, 522] # KD/ODE-init has no sampling-based validation pipeline. - ode = build_training_config( - _job("ode_init", validation_dataset_file="val.json"), "out") + ode = build_training_config(_job("ode_init", validation_dataset_file="val.json"), "out") assert "validation" not in ode["callbacks"] def test_ltx2_models_are_rejected() -> None: assert is_ltx2_model("Lightricks/LTX-2-19B") with pytest.raises(ValueError, match="LTX-2 training is not supported"): - build_training_config( - _job("full_t2v", model_id="Lightricks/LTX-2-19B"), "out") + build_training_config(_job("full_t2v", model_id="Lightricks/LTX-2-19B"), "out") def test_unknown_workload_is_rejected() -> None: @@ -195,8 +190,7 @@ def test_unknown_workload_is_rejected() -> None: def test_invalid_denoising_steps_are_rejected() -> None: with pytest.raises(ValueError, match="Invalid DMD denoising steps"): - build_training_config( - _job("dmd_t2v", dmd_denoising_steps="1000,abc"), "out") + build_training_config(_job("dmd_t2v", dmd_denoising_steps="1000,abc"), "out") def test_training_env_has_no_backend_override() -> None: @@ -211,8 +205,7 @@ def test_workloads_match_frontend_job_config() -> None: (src/lib/jobConfig.ts) — drift means creatable-but-unrunnable jobs.""" import re - job_config = (Path(__file__).resolve().parents[1] / "src" / "lib" / - "jobConfig.ts").read_text(encoding="utf-8") + job_config = (Path(__file__).resolve().parents[1] / "src" / "lib" / "jobConfig.ts").read_text(encoding="utf-8") all_types = set(re.findall(r'type:\s*"([^"]+)"', job_config)) inference_types = {"t2v", "i2v", "t2i"} assert inference_types <= all_types, "jobConfig.ts parse failed" diff --git a/examples/inference/basic/basic.py b/examples/inference/basic/basic.py index 766b9bd79..90a702041 100644 --- a/examples/inference/basic/basic.py +++ b/examples/inference/basic/basic.py @@ -3,6 +3,8 @@ from fastvideo import VideoGenerator # from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. @@ -12,11 +14,11 @@ def main(): "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory + use_fsdp_inference=False, # set to True if GPU is out of memory dit_cpu_offload=False, vae_cpu_offload=False, text_encoder_cpu_offload=True, - pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" + pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" # image_encoder_cpu_offload=False, ) @@ -24,22 +26,19 @@ def main(): # sampling_param.num_frames = 45 # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" # Generate videos with the same simple API, regardless of GPU count - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True) # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/") # Generate another video with a different prompt, without reloading the # model! - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") + prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True) diff --git a/examples/inference/basic/basic_cosmos2_5_i2w.py b/examples/inference/basic/basic_cosmos2_5_i2w.py index 07c5f1316..667bb64fc 100644 --- a/examples/inference/basic/basic_cosmos2_5_i2w.py +++ b/examples/inference/basic/basic_cosmos2_5_i2w.py @@ -30,8 +30,7 @@ def main(): "and casting reflections onto adjacent vehicles. " "The motion creates space in the lineup, signaling activity within the otherwise quiet station. " "It then comes to a smooth stop, resuming its position in line. " - "Overhead signage in Chinese characters remains illuminated, enhancing the vibrant, urban night scene." - ) + "Overhead signage in Chinese characters remains illuminated, enhancing the vibrant, urban night scene.") generator.generate_video( prompt, @@ -47,4 +46,3 @@ def main(): if __name__ == "__main__": main() - diff --git a/examples/inference/basic/basic_cosmos2_5_t2w.py b/examples/inference/basic/basic_cosmos2_5_t2w.py index 2296ca7bd..13590dfdf 100644 --- a/examples/inference/basic/basic_cosmos2_5_t2w.py +++ b/examples/inference/basic/basic_cosmos2_5_t2w.py @@ -31,8 +31,7 @@ def main(): "The welding torch consistently emits sparks and light, and the smoke continues to rise, diffusing slightly as it moves upward. " "The metal surface beneath the torch shows ongoing signs of heating and melting. " "The scene retains its industrial ambiance, with the welding sparks and smoke dominating the visual field, " - "underscoring the ongoing nature of the welding operation." - ) + "underscoring the ongoing nature of the welding operation.") generator.generate_video( prompt, @@ -46,6 +45,3 @@ def main(): if __name__ == "__main__": main() - - - diff --git a/examples/inference/basic/basic_cosmos2_5_v2w.py b/examples/inference/basic/basic_cosmos2_5_v2w.py index 0efefce0e..c174aced9 100644 --- a/examples/inference/basic/basic_cosmos2_5_v2w.py +++ b/examples/inference/basic/basic_cosmos2_5_v2w.py @@ -50,4 +50,3 @@ def main(): if __name__ == "__main__": main() - diff --git a/examples/inference/basic/basic_dmd.py b/examples/inference/basic/basic_dmd.py index 348a403e8..b0340e1c1 100644 --- a/examples/inference/basic/basic_dmd.py +++ b/examples/inference/basic/basic_dmd.py @@ -5,6 +5,8 @@ from fastvideo import VideoGenerator from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_dmd2" + + def main(): os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN" @@ -14,10 +16,10 @@ def main(): model_name, # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory + use_fsdp_inference=False, # set to True if GPU is out of memory # Adjust these offload parameters if you have < 32GB of VRAM text_encoder_cpu_offload=True, - pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" + pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" dit_cpu_offload=False, vae_cpu_offload=False, VSA_sparsity=0.8, @@ -25,7 +27,6 @@ def main(): load_end_time = time.perf_counter() load_time = load_end_time - load_start_time - sampling_param = SamplingParam.from_pretrained(model_name) sampling_param.num_frames = 81 @@ -39,18 +40,16 @@ def main(): # Generate another video with a different prompt, without reloading the # model! - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") + prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") start_time = time.perf_counter() video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, num_frames=81) end_time = time.perf_counter() gen_time2 = end_time - start_time - print(f"Time taken to load model: {load_time} seconds") print(f"Time taken to generate video: {gen_time} seconds") print(f"Time taken to generate video2: {gen_time2} seconds") diff --git a/examples/inference/basic/basic_dmd_new_api.py b/examples/inference/basic/basic_dmd_new_api.py index fe40b32ad..eb319fad4 100644 --- a/examples/inference/basic/basic_dmd_new_api.py +++ b/examples/inference/basic/basic_dmd_new_api.py @@ -32,11 +32,9 @@ def main(): ), # PR 2 still routes a few advanced inference knobs through the # compatibility bridge until they get first-class typed fields. - pipeline=PipelineSelection( - experimental={ - "VSA_sparsity": 0.8, - }, - ), + pipeline=PipelineSelection(experimental={ + "VSA_sparsity": 0.8, + }, ), ) load_start_time = time.perf_counter() @@ -44,14 +42,12 @@ def main(): load_end_time = time.perf_counter() load_time = load_end_time - load_start_time - prompt = ( - "A neon-lit alley in futuristic Tokyo during a heavy rainstorm at night. " - "The puddles reflect glowing signs in kanji, advertising ramen, karaoke, " - "and VR arcades. A woman in a translucent raincoat walks briskly with an " - "LED umbrella. Steam rises from a street food cart, and a cat darts " - "across the screen. Raindrops are visible on the camera lens, creating " - "a cinematic bokeh effect." - ) + prompt = ("A neon-lit alley in futuristic Tokyo during a heavy rainstorm at night. " + "The puddles reflect glowing signs in kanji, advertising ramen, karaoke, " + "and VR arcades. A woman in a translucent raincoat walks briskly with an " + "LED umbrella. Steam rises from a street food cart, and a cat darts " + "across the screen. Raindrops are visible on the camera lens, creating " + "a cinematic bokeh effect.") request = GenerationRequest( prompt=prompt, output=OutputConfig( @@ -66,13 +62,11 @@ def main(): end_time = time.perf_counter() gen_time = end_time - start_time - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently " - "in the breeze, enhancing the lion's commanding presence. The tone is " - "vibrant, embodying the raw energy of the wild. Low angle, steady " - "tracking shot, cinematic." - ) + prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently " + "in the breeze, enhancing the lion's commanding presence. The tone is " + "vibrant, embodying the raw energy of the wild. Low angle, steady " + "tracking shot, cinematic.") request2 = GenerationRequest( prompt=prompt2, output=OutputConfig( diff --git a/examples/inference/basic/basic_dreamx_world.py b/examples/inference/basic/basic_dreamx_world.py index e07f1e1a6..9af4a02ed 100644 --- a/examples/inference/basic/basic_dreamx_world.py +++ b/examples/inference/basic/basic_dreamx_world.py @@ -2,7 +2,6 @@ import os from fastvideo import VideoGenerator - OUTPUT_PATH = os.getenv("DREAMX_WORLD_OUTPUT_PATH", "video_samples_dreamx_world") @@ -46,10 +45,8 @@ def main(): "num_inference_steps": _env_int("DREAMX_WORLD_STEPS", 30), "guidance_scale": _env_float("DREAMX_WORLD_GUIDANCE", 5.0), "action_list": os.getenv("DREAMX_WORLD_ACTIONS", "w,d,w").split(","), - "action_speed_list": [ - float(value) - for value in os.getenv("DREAMX_WORLD_ACTION_SPEEDS", "4.0,2.0,4.0").split(",") - ], + "action_speed_list": + [float(value) for value in os.getenv("DREAMX_WORLD_ACTION_SPEEDS", "4.0,2.0,4.0").split(",")], } if image_path: kwargs["image_path"] = image_path diff --git a/examples/inference/basic/basic_flux2.py b/examples/inference/basic/basic_flux2.py index 8eaeab26d..aa67a2aaf 100644 --- a/examples/inference/basic/basic_flux2.py +++ b/examples/inference/basic/basic_flux2.py @@ -64,12 +64,8 @@ def main() -> None: output = Path(args.output) output.parent.mkdir(parents=True, exist_ok=True) - tp_size = args.tp_size if args.tp_size is not None else ( - args.num_gpus if args.num_gpus > 1 else 1 - ) - sp_size = args.sp_size if args.sp_size is not None else ( - 1 if args.num_gpus > 1 else args.num_gpus - ) + tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1) + sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus) generator_config = GeneratorConfig( model_path=args.model_path, diff --git a/examples/inference/basic/basic_flux2_klein.py b/examples/inference/basic/basic_flux2_klein.py index bc448c6e2..f42525974 100644 --- a/examples/inference/basic/basic_flux2_klein.py +++ b/examples/inference/basic/basic_flux2_klein.py @@ -21,7 +21,6 @@ from fastvideo.api import ( SamplingConfig, ) - DEFAULT_PROMPT = "a brushed steel espresso machine on a marble counter, morning window light" diff --git a/examples/inference/basic/basic_flux_dev.py b/examples/inference/basic/basic_flux_dev.py index c1905128f..930f0f0fb 100644 --- a/examples/inference/basic/basic_flux_dev.py +++ b/examples/inference/basic/basic_flux_dev.py @@ -9,11 +9,9 @@ import re DEFAULT_PROMPTS = [ "a photo of a cat", - ( - "a cinematic photo of a red panda wearing a tiny backpack, standing on a " - "rainy neon-lit street at night, shallow depth of field, sharp focus, " - "35mm, bokeh" - ), + ("a cinematic photo of a red panda wearing a tiny backpack, standing on a " + "rainy neon-lit street at night, shallow depth of field, sharp focus, " + "35mm, bokeh"), ] @@ -42,9 +40,7 @@ def _remove_existing_outputs(out_dir: str, filename_base: str) -> None: def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser( - description="Run FLUX.1-dev text-to-image with FastVideo VideoGenerator.", - ) + p = argparse.ArgumentParser(description="Run FLUX.1-dev text-to-image with FastVideo VideoGenerator.", ) p.add_argument( "--model-path", default="official_weights/FLUX.1-dev", @@ -108,9 +104,7 @@ def main() -> None: try: for i, prompt in enumerate(prompts): seed = args.seed + i - filename_base = ( - f"flux_dev_{i:02d}_seed{seed}_{_safe_filename(prompt, max_len=80)}" - ) + filename_base = (f"flux_dev_{i:02d}_seed{seed}_{_safe_filename(prompt, max_len=80)}") _remove_existing_outputs(args.out_dir, filename_base) output_path = os.path.join(args.out_dir, f"{filename_base}.png") print(f"[flux] prompt_idx={i} seed={seed} output_path={output_path}") diff --git a/examples/inference/basic/basic_gamecraft.py b/examples/inference/basic/basic_gamecraft.py index 17d21fc29..c086d86e4 100644 --- a/examples/inference/basic/basic_gamecraft.py +++ b/examples/inference/basic/basic_gamecraft.py @@ -33,21 +33,20 @@ MODEL_PATH = os.environ.get("GAMECRAFT_MODEL_PATH", "FastVideo/HunyuanGameCraft- # Default prompts for demo DEFAULT_PROMPTS = { - "village": "A charming medieval village with cobblestone streets, thatched-roof houses, and vibrant flower gardens under a bright blue sky.", - "temple": "A majestic ancient temple stands under a clear blue sky, its grandeur highlighted by towering Doric columns and intricate architectural details.", - "forest": "A lush green forest with tall trees, dappled sunlight filtering through the leaves, and a winding dirt path.", + "village": + "A charming medieval village with cobblestone streets, thatched-roof houses, and vibrant flower gardens under a bright blue sky.", + "temple": + "A majestic ancient temple stands under a clear blue sky, its grandeur highlighted by towering Doric columns and intricate architectural details.", + "forest": + "A lush green forest with tall trees, dappled sunlight filtering through the leaves, and a winding dirt path.", "beach": "A tropical beach with crystal clear turquoise water, white sand, and palm trees swaying in the breeze.", } # I2V: default reference image (URL). Can override with a local path. -DEFAULT_I2V_IMAGE_URL = ( - "https://huggingface.co/datasets/huggingface/documentation-images/" - "resolve/main/diffusers/astronaut.jpg" -) -DEFAULT_I2V_PROMPT = ( - "An astronaut hatching from an egg, on the surface of the moon, " - "the darkness and depth of space realised in the background." -) +DEFAULT_I2V_IMAGE_URL = ("https://huggingface.co/datasets/huggingface/documentation-images/" + "resolve/main/diffusers/astronaut.jpg") +DEFAULT_I2V_PROMPT = ("An astronaut hatching from an egg, on the surface of the moon, " + "the darkness and depth of space realised in the background.") OUTPUT_PATH = "video_samples_gamecraft" diff --git a/examples/inference/basic/basic_gen3c.py b/examples/inference/basic/basic_gen3c.py index a91ad301f..d98364b00 100644 --- a/examples/inference/basic/basic_gen3c.py +++ b/examples/inference/basic/basic_gen3c.py @@ -26,51 +26,35 @@ from fastvideo import VideoGenerator def main(): parser = argparse.ArgumentParser(description="GEN3C video generation") - parser.add_argument("--model_path", - type=str, - default="converted_weights/GEN3C-Cosmos-7B") - parser.add_argument("--image_path", - type=str, - default=None, - help="Input image for 3D cache conditioning") - parser.add_argument("--prompt", - type=str, - default="A slow camera pan over a sunlit landscape.") + parser.add_argument("--model_path", type=str, default="converted_weights/GEN3C-Cosmos-7B") + parser.add_argument("--image_path", type=str, default=None, help="Input image for 3D cache conditioning") + parser.add_argument("--prompt", type=str, default="A slow camera pan over a sunlit landscape.") parser.add_argument( "--negative_prompt", type=str, - default=( - "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " - "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " - "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, " - "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special " - "effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and " - "flickering. Overall, the video is of poor quality." - ), + default=("The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " + "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " + "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, " + "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special " + "effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and " + "flickering. Overall, the video is of poor quality."), ) - parser.add_argument("--trajectory", - type=str, - default="left", - choices=[ - "left", "right", "up", "down", "zoom_in", - "zoom_out", "clockwise", "counterclockwise", "none" - ]) + parser.add_argument( + "--trajectory", + type=str, + default="left", + choices=["left", "right", "up", "down", "zoom_in", "zoom_out", "clockwise", "counterclockwise", "none"]) parser.add_argument("--movement_distance", type=float, default=0.3) parser.add_argument("--camera_rotation", type=str, default="center_facing", - choices=[ - "center_facing", "no_rotation", - "trajectory_aligned" - ]) + choices=["center_facing", "no_rotation", "trajectory_aligned"]) parser.add_argument("--height", type=int, default=704) parser.add_argument("--width", type=int, default=1280) parser.add_argument("--num_frames", type=int, default=121) parser.add_argument("--num_inference_steps", type=int, default=35) parser.add_argument("--guidance_scale", type=float, default=1.0) - parser.add_argument("--output_path", - type=str, - default="outputs_video/gen3c.mp4") + parser.add_argument("--output_path", type=str, default="outputs_video/gen3c.mp4") parser.add_argument("--seed", type=int, default=42) args = parser.parse_args() diff --git a/examples/inference/basic/basic_hy15.py b/examples/inference/basic/basic_hy15.py index 745fd7e64..a9dc3d141 100644 --- a/examples/inference/basic/basic_hy15.py +++ b/examples/inference/basic/basic_hy15.py @@ -3,6 +3,8 @@ import json # from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_hy15" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. @@ -12,31 +14,38 @@ def main(): "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory + use_fsdp_inference=False, # set to True if GPU is out of memory dit_cpu_offload=True, vae_cpu_offload=True, text_encoder_cpu_offload=True, - pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" + pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" # image_encoder_cpu_offload=False, ) - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") - video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=81, fps=16) + video = generator.generate_video(prompt, + output_path=OUTPUT_PATH, + save_video=True, + negative_prompt="", + num_frames=81, + fps=16) - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") + prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") - video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=81, fps=16) + video2 = generator.generate_video(prompt2, + output_path=OUTPUT_PATH, + save_video=True, + negative_prompt="", + num_frames=81, + fps=16) if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/inference/basic/basic_hy15_1080p.py b/examples/inference/basic/basic_hy15_1080p.py index 32dcc5780..29ce70152 100644 --- a/examples/inference/basic/basic_hy15_1080p.py +++ b/examples/inference/basic/basic_hy15_1080p.py @@ -3,38 +3,37 @@ import json # from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_hy15_1080p" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. # If a local path is provided, FastVideo will make a best effort # attempt to identify the optimal arguments. generator = VideoGenerator.from_pretrained( - "weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR", # 480p -> 720p -> 1080p + "weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR", # 480p -> 720p -> 1080p # or "weizhou03/HunyuanVideo-1.5-Diffusers-1080p" # 720p -> 1080p # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory + use_fsdp_inference=False, # set to True if GPU is out of memory dit_cpu_offload=True, vae_cpu_offload=True, text_encoder_cpu_offload=True, - pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" + pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" # image_encoder_cpu_offload=False, ) - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="") - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") + prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, negative_prompt="") diff --git a/examples/inference/basic/basic_hyworld.py b/examples/inference/basic/basic_hyworld.py index 4caae8656..72bcf6593 100644 --- a/examples/inference/basic/basic_hyworld.py +++ b/examples/inference/basic/basic_hyworld.py @@ -6,6 +6,8 @@ DEFAULT_PROMPT = 'A paved pathway leads towards a stone arch bridge spanning a c DEFAULT_IMAGE = 'https://raw.githubusercontent.com/Tencent-Hunyuan/HY-WorldPlay/main/assets/img/test.png' OUTPUT_PATH = "video_samples_hyworld" + + def main(): import argparse diff --git a/examples/inference/basic/basic_kandinsky5_i2v.py b/examples/inference/basic/basic_kandinsky5_i2v.py index c30ca2f48..a11dde536 100644 --- a/examples/inference/basic/basic_kandinsky5_i2v.py +++ b/examples/inference/basic/basic_kandinsky5_i2v.py @@ -7,7 +7,7 @@ IMAGE_PATH = "assets/girl.png" def main(): generator = VideoGenerator.from_pretrained( - "kandinskylab/Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers", + "kandinskylab/Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers", # "kandinskylab/Kandinsky-5.0-I2V-Pro-sft-5s-Diffusers" # "kandinskylab/Kandinsky-5.0-I2V-Lite-5s-Diffusers" num_gpus=1, @@ -19,9 +19,7 @@ def main(): # image_encoder_cpu_offload=False, ) - prompt = ( - "A woman stands up and walks away" - ) + prompt = ("A woman stands up and walks away") _ = generator.generate_video( prompt, image_path=IMAGE_PATH, diff --git a/examples/inference/basic/basic_kandinsky5_t2v.py b/examples/inference/basic/basic_kandinsky5_t2v.py index 8042f87fd..37ee149a0 100644 --- a/examples/inference/basic/basic_kandinsky5_t2v.py +++ b/examples/inference/basic/basic_kandinsky5_t2v.py @@ -2,6 +2,7 @@ from fastvideo import VideoGenerator OUTPUT_PATH = "video_samples_kandinsky5_t2v" + def main(): generator = VideoGenerator.from_pretrained( "kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers", @@ -17,21 +18,28 @@ def main(): # image_encoder_cpu_offload=False, ) - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) - _ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True,height=512, width=768, num_frames=121) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") + _ = generator.generate_video(prompt, + output_path=OUTPUT_PATH, + save_video=True, + height=512, + width=768, + num_frames=121) - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") - _ = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, height=512, width=768, num_frames=121) + prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") + _ = generator.generate_video(prompt2, + output_path=OUTPUT_PATH, + save_video=True, + height=512, + width=768, + num_frames=121) if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/inference/basic/basic_lingbotworld2_causal_fast.py b/examples/inference/basic/basic_lingbotworld2_causal_fast.py index bea4637d7..4c08764a7 100644 --- a/examples/inference/basic/basic_lingbotworld2_causal_fast.py +++ b/examples/inference/basic/basic_lingbotworld2_causal_fast.py @@ -6,7 +6,6 @@ from pathlib import Path from fastvideo import VideoGenerator - REPO_ROOT = Path(__file__).resolve().parents[3] DATASET_DIR = REPO_ROOT / "examples" / "dataset" / "lingbotworld2" OUTPUT_PATH = REPO_ROOT / "outputs" / "lingbotworld2_causal_fast.mp4" diff --git a/examples/inference/basic/basic_lingbotworld_base_cam.py b/examples/inference/basic/basic_lingbotworld_base_cam.py index 8c9a95321..2be36e022 100644 --- a/examples/inference/basic/basic_lingbotworld_base_cam.py +++ b/examples/inference/basic/basic_lingbotworld_base_cam.py @@ -3,17 +3,19 @@ from fastvideo.models.dits.lingbotworld.cam_utils import prepare_camera_embeddin # from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_lingbotworld" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. # If a local path is provided, FastVideo will make a best effort # attempt to identify the optimal arguments. generator = VideoGenerator.from_pretrained( - "FastVideo/LingBot-World-Base-Cam-Diffusers", + "FastVideo/LingBot-World-Base-Cam-Diffusers", # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory - dit_cpu_offload=True, # DiT need to be offloaded for MoE + use_fsdp_inference=False, # set to True if GPU is out of memory + dit_cpu_offload=True, # DiT need to be offloaded for MoE vae_cpu_offload=False, text_encoder_cpu_offload=True, # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer diff --git a/examples/inference/basic/basic_longcat_i2v.py b/examples/inference/basic/basic_longcat_i2v.py index f170dc4a3..001885247 100644 --- a/examples/inference/basic/basic_longcat_i2v.py +++ b/examples/inference/basic/basic_longcat_i2v.py @@ -21,20 +21,16 @@ import os from fastvideo import VideoGenerator # Common prompts and settings matching the shell script examples -PROMPT = ( - "A woman sits at a wooden table by the window in a cozy café. She reaches out " - "with her right hand, picks up the white coffee cup from the saucer, and gently " - "brings it to her lips to take a sip. After drinking, she places the cup back on " - "the table and looks out the window, enjoying the peaceful atmosphere." -) +PROMPT = ("A woman sits at a wooden table by the window in a cozy café. She reaches out " + "with her right hand, picks up the white coffee cup from the saucer, and gently " + "brings it to her lips to take a sip. After drinking, she places the cup back on " + "the table and looks out the window, enjoying the peaceful atmosphere.") -NEGATIVE_PROMPT = ( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, " - "paintings, images, static, overall gray, worst quality, low quality, JPEG compression " - "residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, " - "deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, " - "three legs, many people in the background, walking backwards" -) +NEGATIVE_PROMPT = ("Bright tones, overexposed, static, blurred details, subtitles, style, works, " + "paintings, images, static, overall gray, worst quality, low quality, JPEG compression " + "residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, " + "deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, " + "three legs, many people in the background, walking backwards") # Input image path IMAGE_PATH = "assets/girl.png" @@ -51,20 +47,20 @@ def basic_generation(): print("=" * 60) print("LongCat I2V: Basic Generation (50 steps, 480p)") print("=" * 60) - + generator = VideoGenerator.from_pretrained( "FastVideo/LongCat-Video-I2V-Diffusers", num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory + use_fsdp_inference=False, # set to True if GPU is out of memory dit_cpu_offload=False, vae_cpu_offload=True, text_encoder_cpu_offload=True, pin_cpu_memory=False, enable_bsa=False, ) - + output_path = "outputs_video/longcat_i2v_basic" - + generator.generate_video( prompt=PROMPT, negative_prompt=NEGATIVE_PROMPT, @@ -79,7 +75,7 @@ def basic_generation(): guidance_scale=4.0, seed=SEED, ) - + print(f"\nBasic generation complete! Video saved to: {output_path}") generator.shutdown() @@ -94,11 +90,11 @@ def distill_refine_generation(): print("\n" + "=" * 60) print("LongCat I2V: Distill + Refine Pipeline") print("=" * 60) - + # Stage 1: Distilled generation (16 steps at 480p) print("\n[Stage 1] Distilled generation (16 steps, 480p)") print("-" * 40) - + generator = VideoGenerator.from_pretrained( "FastVideo/LongCat-Video-I2V-Diffusers", num_gpus=1, @@ -111,9 +107,9 @@ def distill_refine_generation(): lora_path="FastVideo/LongCat-Video-T2V-Distilled-LoRA", lora_nickname="distilled", ) - + distill_output_path = "outputs_video/longcat_i2v_distill" - + generator.generate_video( prompt=PROMPT, negative_prompt=NEGATIVE_PROMPT, @@ -128,14 +124,14 @@ def distill_refine_generation(): guidance_scale=1.0, seed=SEED, ) - + print(f"Distilled generation complete! Video saved to: {distill_output_path}") generator.shutdown() - + # Stage 2: Refinement (480p -> 768p) print("\n[Stage 2] Refinement (480p -> 768p with BSA)") print("-" * 40) - + # Find the actual saved video file from stage 1 video_files = glob.glob(os.path.join(distill_output_path, "*.mp4")) if not video_files: @@ -143,7 +139,7 @@ def distill_refine_generation(): # Use the most recently created video file distill_video_path = max(video_files, key=os.path.getmtime) print(f"Using stage 1 video: {distill_video_path}") - + # Create a new generator with refinement LoRA and BSA enabled # Note: Refinement uses the T2V model (not I2V) since it's upscaling the generated video # For BSA [4, 4, 8]: latent must be divisible by 8 @@ -163,9 +159,9 @@ def distill_refine_generation(): lora_path="FastVideo/LongCat-Video-T2V-Refinement-LoRA", lora_nickname="refinement", ) - + refine_output_path = "outputs_video/longcat_i2v_refine_720p" - + refine_generator.generate_video( prompt=PROMPT, negative_prompt=NEGATIVE_PROMPT, @@ -182,7 +178,7 @@ def distill_refine_generation(): guidance_scale=1.0, seed=SEED, ) - + print(f"Refinement complete! Video saved to: {refine_output_path}") refine_generator.shutdown() @@ -192,13 +188,13 @@ def main(): print("\n" + "=" * 60) print("LongCat Image-to-Video Example") print("=" * 60 + "\n") - + # Run basic generation basic_generation() - + # Run distill+refine pipeline distill_refine_generation() - + print("\n" + "=" * 60) print("All generations complete!") print("=" * 60) @@ -206,5 +202,3 @@ def main(): if __name__ == "__main__": main() - - diff --git a/examples/inference/basic/basic_longcat_t2v.py b/examples/inference/basic/basic_longcat_t2v.py index 354228de2..58c87dc3a 100644 --- a/examples/inference/basic/basic_longcat_t2v.py +++ b/examples/inference/basic/basic_longcat_t2v.py @@ -15,22 +15,18 @@ import os from fastvideo import VideoGenerator # Common prompts and settings matching the shell script examples -PROMPT = ( - "In a realistic photography style, a white boy around seven or eight years old " - "sits on a park bench, wearing a light blue T-shirt, denim shorts, and white sneakers. " - "He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a " - "medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, " - "who eagerly licks it with its tongue. The sun is shining brightly, and the background " - "features a green lawn and several tall trees, creating a warm and loving scene." -) +PROMPT = ("In a realistic photography style, a white boy around seven or eight years old " + "sits on a park bench, wearing a light blue T-shirt, denim shorts, and white sneakers. " + "He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a " + "medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, " + "who eagerly licks it with its tongue. The sun is shining brightly, and the background " + "features a green lawn and several tall trees, creating a warm and loving scene.") -NEGATIVE_PROMPT = ( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, " - "paintings, images, static, overall gray, worst quality, low quality, JPEG compression " - "residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, " - "deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, " - "three legs, many people in the background, walking backwards" -) +NEGATIVE_PROMPT = ("Bright tones, overexposed, static, blurred details, subtitles, style, works, " + "paintings, images, static, overall gray, worst quality, low quality, JPEG compression " + "residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, " + "deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, " + "three legs, many people in the background, walking backwards") SEED = 42 @@ -44,20 +40,20 @@ def basic_generation(): print("=" * 60) print("LongCat T2V: Basic Generation (50 steps, 480p)") print("=" * 60) - + generator = VideoGenerator.from_pretrained( "FastVideo/LongCat-Video-T2V-Diffusers", num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory + use_fsdp_inference=False, # set to True if GPU is out of memory dit_cpu_offload=False, vae_cpu_offload=True, text_encoder_cpu_offload=True, pin_cpu_memory=False, enable_bsa=False, ) - + output_path = "outputs_video/longcat_t2v_basic" - + generator.generate_video( prompt=PROMPT, negative_prompt=NEGATIVE_PROMPT, @@ -71,7 +67,7 @@ def basic_generation(): guidance_scale=4.0, seed=SEED, ) - + print(f"\nBasic generation complete! Video saved to: {output_path}") generator.shutdown() @@ -86,11 +82,11 @@ def distill_refine_generation(): print("\n" + "=" * 60) print("LongCat T2V: Distill + Refine Pipeline") print("=" * 60) - + # Stage 1: Distilled generation (16 steps at 480p) print("\n[Stage 1] Distilled generation (16 steps, 480p)") print("-" * 40) - + generator = VideoGenerator.from_pretrained( "FastVideo/LongCat-Video-T2V-Diffusers", num_gpus=1, @@ -103,9 +99,9 @@ def distill_refine_generation(): lora_path="FastVideo/LongCat-Video-T2V-Distilled-LoRA", lora_nickname="distilled", ) - + distill_output_path = "outputs_video/longcat_t2v_distill" - + generator.generate_video( prompt=PROMPT, negative_prompt=NEGATIVE_PROMPT, @@ -119,14 +115,14 @@ def distill_refine_generation(): guidance_scale=1.0, seed=SEED, ) - + print(f"Distilled generation complete! Video saved to: {distill_output_path}") generator.shutdown() - + # Stage 2: Refinement (480p -> 720p) print("\n[Stage 2] Refinement (480p -> 720p with BSA)") print("-" * 40) - + # Find the actual saved video file from stage 1 video_files = glob.glob(os.path.join(distill_output_path, "*.mp4")) if not video_files: @@ -134,7 +130,7 @@ def distill_refine_generation(): # Use the most recently created video file distill_video_path = max(video_files, key=os.path.getmtime) print(f"Using stage 1 video: {distill_video_path}") - + # Create a new generator with refinement LoRA and BSA enabled refine_generator = VideoGenerator.from_pretrained( "FastVideo/LongCat-Video-T2V-Diffusers", @@ -151,9 +147,9 @@ def distill_refine_generation(): lora_path="FastVideo/LongCat-Video-T2V-Refinement-LoRA", lora_nickname="refinement", ) - + refine_output_path = "outputs_video/longcat_t2v_refine_720p" - + refine_generator.generate_video( prompt=PROMPT, negative_prompt=NEGATIVE_PROMPT, @@ -170,7 +166,7 @@ def distill_refine_generation(): guidance_scale=1.0, seed=SEED, ) - + print(f"Refinement complete! Video saved to: {refine_output_path}") refine_generator.shutdown() @@ -180,13 +176,13 @@ def main(): print("\n" + "=" * 60) print("LongCat Text-to-Video Example") print("=" * 60 + "\n") - + # Run basic generation basic_generation() - + # Run distill+refine pipeline distill_refine_generation() - + print("\n" + "=" * 60) print("All generations complete!") print("=" * 60) @@ -194,5 +190,3 @@ def main(): if __name__ == "__main__": main() - - diff --git a/examples/inference/basic/basic_longcat_vc.py b/examples/inference/basic/basic_longcat_vc.py index 313cbb72c..36d5f0935 100644 --- a/examples/inference/basic/basic_longcat_vc.py +++ b/examples/inference/basic/basic_longcat_vc.py @@ -21,21 +21,17 @@ import os from fastvideo import VideoGenerator # Common prompts and settings matching the shell script examples -PROMPT = ( - "A person rides a motorcycle along a long, straight road that stretches between " - "a body of water and a forested hillside. The rider steadily accelerates, keeping " - "the motorcycle centered between the guardrails, while the scenery passes by on " - "both sides. The video captures the journey from the rider's perspective, emphasizing " - "the sense of motion and adventure." -) +PROMPT = ("A person rides a motorcycle along a long, straight road that stretches between " + "a body of water and a forested hillside. The rider steadily accelerates, keeping " + "the motorcycle centered between the guardrails, while the scenery passes by on " + "both sides. The video captures the journey from the rider's perspective, emphasizing " + "the sense of motion and adventure.") -NEGATIVE_PROMPT = ( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, " - "paintings, images, static, overall gray, worst quality, low quality, JPEG compression " - "residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, " - "deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, " - "three legs, many people in the background, walking backwards" -) +NEGATIVE_PROMPT = ("Bright tones, overexposed, static, blurred details, subtitles, style, works, " + "paintings, images, static, overall gray, worst quality, low quality, JPEG compression " + "residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, " + "deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, " + "three legs, many people in the background, walking backwards") # Input video path VIDEO_PATH = "assets/motorcycle.mp4" @@ -55,27 +51,25 @@ def basic_generation(): print("=" * 60) print("LongCat VC: Basic Generation (50 steps, 480p)") print("=" * 60) - + # Check if video exists if not os.path.exists(VIDEO_PATH): - raise FileNotFoundError( - f"Video not found at {VIDEO_PATH}. " - "Please provide a valid video path." - ) - + raise FileNotFoundError(f"Video not found at {VIDEO_PATH}. " + "Please provide a valid video path.") + generator = VideoGenerator.from_pretrained( "FastVideo/LongCat-Video-VC-Diffusers", num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory + use_fsdp_inference=False, # set to True if GPU is out of memory dit_cpu_offload=False, vae_cpu_offload=True, text_encoder_cpu_offload=True, pin_cpu_memory=False, enable_bsa=False, ) - + output_path = "outputs_video/longcat_vc_basic" - + generator.generate_video( prompt=PROMPT, negative_prompt=NEGATIVE_PROMPT, @@ -91,7 +85,7 @@ def basic_generation(): guidance_scale=4.0, seed=SEED, ) - + print(f"\nBasic generation complete! Video saved to: {output_path}") generator.shutdown() @@ -106,18 +100,16 @@ def distill_refine_generation(): print("\n" + "=" * 60) print("LongCat VC: Distill + Refine Pipeline") print("=" * 60) - + # Check if video exists if not os.path.exists(VIDEO_PATH): - raise FileNotFoundError( - f"Video not found at {VIDEO_PATH}. " - "Please provide a valid video path." - ) - + raise FileNotFoundError(f"Video not found at {VIDEO_PATH}. " + "Please provide a valid video path.") + # Stage 1: Distilled generation (16 steps at 480p) print("\n[Stage 1] Distilled generation (16 steps, 480p)") print("-" * 40) - + generator = VideoGenerator.from_pretrained( "FastVideo/LongCat-Video-VC-Diffusers", num_gpus=1, @@ -130,9 +122,9 @@ def distill_refine_generation(): lora_path="FastVideo/LongCat-Video-T2V-Distilled-LoRA", lora_nickname="distilled", ) - + distill_output_path = "outputs_video/longcat_vc_distill" - + generator.generate_video( prompt=PROMPT, negative_prompt=NEGATIVE_PROMPT, @@ -148,14 +140,14 @@ def distill_refine_generation(): guidance_scale=1.0, seed=SEED, ) - + print(f"Distilled generation complete! Video saved to: {distill_output_path}") generator.shutdown() - + # Stage 2: Refinement (480p -> 720p) print("\n[Stage 2] Refinement (480p -> 720p with BSA)") print("-" * 40) - + # Find the actual saved video file from stage 1 video_files = glob.glob(os.path.join(distill_output_path, "*.mp4")) if not video_files: @@ -163,7 +155,7 @@ def distill_refine_generation(): # Use the most recently created video file distill_video_path = max(video_files, key=os.path.getmtime) print(f"Using stage 1 video: {distill_video_path}") - + # Create a new generator with refinement LoRA and BSA enabled # Note: Refinement uses the T2V model (not VC) since it's upscaling the generated video refine_generator = VideoGenerator.from_pretrained( @@ -181,9 +173,9 @@ def distill_refine_generation(): lora_path="FastVideo/LongCat-Video-T2V-Refinement-LoRA", lora_nickname="refinement", ) - + refine_output_path = "outputs_video/longcat_vc_refine_720p" - + refine_generator.generate_video( prompt=PROMPT, negative_prompt=NEGATIVE_PROMPT, @@ -200,7 +192,7 @@ def distill_refine_generation(): guidance_scale=1.0, seed=SEED, ) - + print(f"Refinement complete! Video saved to: {refine_output_path}") refine_generator.shutdown() @@ -210,13 +202,13 @@ def main(): print("\n" + "=" * 60) print("LongCat Video Continuation Example") print("=" * 60 + "\n") - + # Run basic generation basic_generation() - + # Run distill+refine pipeline distill_refine_generation() - + print("\n" + "=" * 60) print("All generations complete!") print("=" * 60) @@ -224,5 +216,3 @@ def main(): if __name__ == "__main__": main() - - diff --git a/examples/inference/basic/basic_ltx2.py b/examples/inference/basic/basic_ltx2.py index c7196e291..e7e1dea00 100644 --- a/examples/inference/basic/basic_ltx2.py +++ b/examples/inference/basic/basic_ltx2.py @@ -1,19 +1,16 @@ from fastvideo import VideoGenerator - -PROMPT = ( - "A warm sunny backyard. The camera starts in a tight cinematic close-up " - "of a woman and a man in their 30s, facing each other with serious " - "expressions. The woman, emotional and dramatic, says softly, \"That's " - "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " - "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " - "then mutters defensively, \"He's just having fun.\" The camera slowly " - "pans right, revealing the grandfather in the garden wearing enormous " - "butterfly wings, waving his arms in the air like he's trying to take " - "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. " - "The woman covers her face, on the verge of tears. The tone is deadpan, " - "absurd, and quietly tragic." -) +PROMPT = ("A warm sunny backyard. The camera starts in a tight cinematic close-up " + "of a woman and a man in their 30s, facing each other with serious " + "expressions. The woman, emotional and dramatic, says softly, \"That's " + "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " + "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " + "then mutters defensively, \"He's just having fun.\" The camera slowly " + "pans right, revealing the grandfather in the garden wearing enormous " + "butterfly wings, waving his arms in the air like he's trying to take " + "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. " + "The woman covers her face, on the verge of tears. The tone is deadpan, " + "absurd, and quietly tragic.") def main() -> None: @@ -36,4 +33,4 @@ def main() -> None: if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/inference/basic/basic_ltx2_3_distilled_i2v.py b/examples/inference/basic/basic_ltx2_3_distilled_i2v.py index 7e680d032..8026bf2cb 100644 --- a/examples/inference/basic/basic_ltx2_3_distilled_i2v.py +++ b/examples/inference/basic/basic_ltx2_3_distilled_i2v.py @@ -67,25 +67,19 @@ _inductor.coordinate_descent_tuning = True _inductor.coordinate_descent_check_all_directions = True _inductor.epilogue_fusion = False -MODEL_ID = os.path.expandvars( - os.path.expanduser( - os.getenv("LTX23_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers") - ) -) -OUTPUT_DIR = Path( - os.getenv("LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_distilled_i2v") -) +MODEL_ID = os.path.expandvars(os.path.expanduser(os.getenv("LTX23_MODEL_PATH", + "FastVideo/LTX-2.3-Distilled-Diffusers"))) +OUTPUT_DIR = Path(os.getenv("LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_distilled_i2v")) I2V_IMAGE = os.getenv("LTX23_I2V_IMAGE", "") -DEFAULT_PROMPT = ( - "A fashion model takes a slow step forward and shifts her weight, " - "the soft fabric of her clothing swaying and rippling with the " - "motion, her hair shifting gently, soft even studio lighting on a " - "clean light background, elegant slow-motion runway feel." -) +DEFAULT_PROMPT = ("A fashion model takes a slow step forward and shifts her weight, " + "the soft fabric of her clothing swaying and rippling with the " + "motion, her hair shifting gently, soft even studio lighting on a " + "clean light background, elegant slow-motion runway feel.") PROMPT = os.getenv("LTX23_I2V_PROMPT", DEFAULT_PROMPT) # Per-stage timing helpers -------------------------------------------------- + def _print_stage_breakdown(result: dict, label: str) -> float | None: """Print stage execution times and return the sum, or None if missing.""" logging_info = result.get("logging_info") @@ -114,9 +108,7 @@ def _collect_stage_times( return for name, metrics in stages.items(): stage_order.setdefault(name, None) - stage_times.setdefault(name, []).append( - float(metrics.get("execution_time", 0.0)) - ) + stage_times.setdefault(name, []).append(float(metrics.get("execution_time", 0.0))) def _resolve_refine_upsampler(model_root: str) -> Path: @@ -125,21 +117,18 @@ def _resolve_refine_upsampler(model_root: str) -> Path: cand = Path(model_root) / name if (cand / "config.json").is_file(): return cand - raise FileNotFoundError( - f"No refine upsampler directory under {model_root}. " - f"Expected `{model_root}/spatial_upscaler/config.json`." - ) + raise FileNotFoundError(f"No refine upsampler directory under {model_root}. " + f"Expected `{model_root}/spatial_upscaler/config.json`.") # Main --------------------------------------------------------------------- + def main() -> None: if not I2V_IMAGE: - raise SystemExit( - "LTX23_I2V_IMAGE is required for i2v. Example:\n" - " export LTX23_I2V_IMAGE=/path/to/portrait_or_product.jpg\n" - " python examples/inference/basic/basic_ltx2_3_distilled_i2v.py" - ) + raise SystemExit("LTX23_I2V_IMAGE is required for i2v. Example:\n" + " export LTX23_I2V_IMAGE=/path/to/portrait_or_product.jpg\n" + " python examples/inference/basic/basic_ltx2_3_distilled_i2v.py") if not Path(I2V_IMAGE).is_file(): raise SystemExit(f"LTX23_I2V_IMAGE not found: {I2V_IMAGE}") @@ -201,11 +190,13 @@ def main() -> None: common_kwargs = dict( prompt=PROMPT, - negative_prompt="", # distilled is CFG-free; no negative needed - guidance_scale=1.0, # CFG=1 for distilled - height=1280, width=832, # portrait runway aspect - num_frames=121, fps=24, # ~5s clip - num_inference_steps=8, # distilled denoise steps + negative_prompt="", # distilled is CFG-free; no negative needed + guidance_scale=1.0, # CFG=1 for distilled + height=1280, + width=832, # portrait runway aspect + num_frames=121, + fps=24, # ~5s clip + num_inference_steps=8, # distilled denoise steps # i2v: anchor the input image at frame 0 with full strength. # `ltx2_image_crf=0.0` skips an extra JPEG re-encode of an already # JPEG conditioning image. @@ -251,10 +242,7 @@ def main() -> None: **common_kwargs, ) wall = time.perf_counter() - t0 - e2e = ( - result.get("e2e_latency") - if isinstance(result, dict) else None - ) or wall + e2e = (result.get("e2e_latency") if isinstance(result, dict) else None) or wall measured_secs.append(e2e) print(f"[measured {m + 1}/{measured_runs}] e2e={e2e:.2f}s wall={wall:.2f}s") if isinstance(result, dict): @@ -266,10 +254,8 @@ def main() -> None: print(f"warmup wall-times: {[round(x, 1) for x in warmup_secs]}") if measured_secs: avg = sum(measured_secs) / len(measured_secs) - print( - f"measured e2e (n={len(measured_secs)}): " - f"{[round(x, 2) for x in measured_secs]} -> avg {avg:.2f}s" - ) + print(f"measured e2e (n={len(measured_secs)}): " + f"{[round(x, 2) for x in measured_secs]} -> avg {avg:.2f}s") if stage_times: print(f"average stage times over {measured_runs} measured runs:") avg_total = 0.0 diff --git a/examples/inference/basic/basic_ltx2_3_distilled_i2v_typed.py b/examples/inference/basic/basic_ltx2_3_distilled_i2v_typed.py index defa4fdbb..f5ac88e0f 100644 --- a/examples/inference/basic/basic_ltx2_3_distilled_i2v_typed.py +++ b/examples/inference/basic/basic_ltx2_3_distilled_i2v_typed.py @@ -100,23 +100,14 @@ _inductor.coordinate_descent_tuning = True _inductor.coordinate_descent_check_all_directions = True _inductor.epilogue_fusion = False -MODEL_ID = os.path.expandvars( - os.path.expanduser( - os.getenv("LTX23_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers") - ) -) -OUTPUT_DIR = Path( - os.getenv( - "LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_distilled_i2v_typed" - ) -) +MODEL_ID = os.path.expandvars(os.path.expanduser(os.getenv("LTX23_MODEL_PATH", + "FastVideo/LTX-2.3-Distilled-Diffusers"))) +OUTPUT_DIR = Path(os.getenv("LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_distilled_i2v_typed")) I2V_IMAGE = os.getenv("LTX23_I2V_IMAGE", "") -DEFAULT_PROMPT = ( - "A fashion model takes a slow step forward and shifts her weight, " - "the soft fabric of her clothing swaying and rippling with the " - "motion, her hair shifting gently, soft even studio lighting on a " - "clean light background, elegant slow-motion runway feel." -) +DEFAULT_PROMPT = ("A fashion model takes a slow step forward and shifts her weight, " + "the soft fabric of her clothing swaying and rippling with the " + "motion, her hair shifting gently, soft even studio lighting on a " + "clean light background, elegant slow-motion runway feel.") PROMPT = os.getenv("LTX23_I2V_PROMPT", DEFAULT_PROMPT) @@ -147,9 +138,7 @@ def _collect_stage_times( return for name, metrics in stages.items(): stage_order.setdefault(name, None) - stage_times.setdefault(name, []).append( - float(metrics.get("execution_time", 0.0)) - ) + stage_times.setdefault(name, []).append(float(metrics.get("execution_time", 0.0))) def _resolve_refine_upsampler(model_root: str) -> Path: @@ -157,20 +146,16 @@ def _resolve_refine_upsampler(model_root: str) -> Path: cand = Path(model_root) / name if (cand / "config.json").is_file(): return cand - raise FileNotFoundError( - f"No refine upsampler directory under {model_root}. " - f"Expected `{model_root}/spatial_upscaler/config.json`." - ) + raise FileNotFoundError(f"No refine upsampler directory under {model_root}. " + f"Expected `{model_root}/spatial_upscaler/config.json`.") def main() -> None: if not I2V_IMAGE: - raise SystemExit( - "LTX23_I2V_IMAGE is required for i2v. Example:\n" - " export LTX23_I2V_IMAGE=/path/to/portrait_or_product.jpg\n" - " python examples/inference/basic/" - "basic_ltx2_3_distilled_i2v_typed.py" - ) + raise SystemExit("LTX23_I2V_IMAGE is required for i2v. Example:\n" + " export LTX23_I2V_IMAGE=/path/to/portrait_or_product.jpg\n" + " python examples/inference/basic/" + "basic_ltx2_3_distilled_i2v_typed.py") if not Path(I2V_IMAGE).is_file(): raise SystemExit(f"LTX23_I2V_IMAGE not found: {I2V_IMAGE}") @@ -220,10 +205,9 @@ def main() -> None: # model-specific VAE precision / decoder defaults are picked up # the same way the legacy example's # ``PipelineConfig.from_pretrained(model_root)`` did them. - components=ComponentConfig( - upsampler_weights=str(refine_upsampler_path), - # Distilled has no refine LoRA — omit ``lora_path``. - ), + components=ComponentConfig(upsampler_weights=str(refine_upsampler_path), + # Distilled has no refine LoRA — omit ``lora_path``. + ), vae_tiling=False, preset_overrides={ "refine": { @@ -278,11 +262,7 @@ def main() -> None: for w in range(warmup_runs): print(f"\n[warmup {w + 1}/{warmup_runs}] compiling + generating…") t0 = time.perf_counter() - generator.generate( - build_request( - OUTPUT_DIR / f"_warmup_{w + 1}.mp4", seed=7 - ) - ) + generator.generate(build_request(OUTPUT_DIR / f"_warmup_{w + 1}.mp4", seed=7)) dt = time.perf_counter() - t0 warmup_secs.append(dt) print(f"[warmup {w + 1}/{warmup_runs}] wall={dt:.1f}s") @@ -291,46 +271,30 @@ def main() -> None: (OUTPUT_DIR / f"_warmup_{w + 1}.mp4").unlink(missing_ok=True) for m in range(measured_runs): - out_path = ( - OUTPUT_DIR - / f"output_ltx2_3_distilled_i2v_typed_run_{m + 1}.mp4" - ) - print( - f"\n[measured {m + 1}/{measured_runs}] generating: {out_path}" - ) + out_path = (OUTPUT_DIR / f"output_ltx2_3_distilled_i2v_typed_run_{m + 1}.mp4") + print(f"\n[measured {m + 1}/{measured_runs}] generating: {out_path}") t0 = time.perf_counter() - result = generator.generate( - build_request(out_path, seed=2002 + m) - ) + result = generator.generate(build_request(out_path, seed=2002 + m)) wall = time.perf_counter() - t0 # ``e2e_latency`` is currently surfaced via ``result.extra``; # ``GenerationResult`` exposes ``generation_time`` as a # first-class field but the LTX-2 pipeline only fills the # legacy ``e2e_latency`` key. Prefer the explicit one, fall # back to wall-clock. - e2e = ( - result.extra.get("e2e_latency") - if hasattr(result, "extra") else None - ) or wall + e2e = (result.extra.get("e2e_latency") if hasattr(result, "extra") else None) or wall measured_secs.append(e2e) - print( - f"[measured {m + 1}/{measured_runs}] " - f"e2e={e2e:.2f}s wall={wall:.2f}s" - ) + print(f"[measured {m + 1}/{measured_runs}] " + f"e2e={e2e:.2f}s wall={wall:.2f}s") _print_stage_breakdown(result, f"measured {m + 1}") _collect_stage_times(result, stage_times, stage_order) print("\n=== summary ===") - print( - f"warmup wall-times: " - f"{[round(x, 1) for x in warmup_secs]}" - ) + print(f"warmup wall-times: " + f"{[round(x, 1) for x in warmup_secs]}") if measured_secs: avg = sum(measured_secs) / len(measured_secs) - print( - f"measured e2e (n={len(measured_secs)}): " - f"{[round(x, 2) for x in measured_secs]} -> avg {avg:.2f}s" - ) + print(f"measured e2e (n={len(measured_secs)}): " + f"{[round(x, 2) for x in measured_secs]} -> avg {avg:.2f}s") if stage_times: print(f"average stage times over {measured_runs} measured runs:") avg_total = 0.0 diff --git a/examples/inference/basic/basic_ltx2_distilled.py b/examples/inference/basic/basic_ltx2_distilled.py index 1ac732eaf..1eb1f3c1d 100644 --- a/examples/inference/basic/basic_ltx2_distilled.py +++ b/examples/inference/basic/basic_ltx2_distilled.py @@ -1,21 +1,21 @@ from fastvideo import VideoGenerator -PROMPT = ( - "A warm sunny backyard. The camera starts in a tight cinematic close-up " - "of a woman and a man in their 30s, facing each other with serious " - "expressions. The woman, emotional and dramatic, says softly, \"That's " - "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " - "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " - "then mutters defensively, \"He's just having fun.\" The camera slowly " - "pans right, revealing the grandfather in the garden wearing enormous " - "butterfly wings, waving his arms in the air like he's trying to take " - "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. " - "The woman covers her face, on the verge of tears. The tone is deadpan, " - "absurd, and quietly tragic." -) +PROMPT = ("A warm sunny backyard. The camera starts in a tight cinematic close-up " + "of a woman and a man in their 30s, facing each other with serious " + "expressions. The woman, emotional and dramatic, says softly, \"That's " + "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " + "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " + "then mutters defensively, \"He's just having fun.\" The camera slowly " + "pans right, revealing the grandfather in the garden wearing enormous " + "butterfly wings, waving his arms in the air like he's trying to take " + "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. " + "The woman covers her face, on the verge of tears. The tone is deadpan, " + "absurd, and quietly tragic.") import os + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" + def main() -> None: generator = VideoGenerator.from_pretrained( "FastVideo/LTX2-Distilled-Diffusers", diff --git a/examples/inference/basic/basic_ltx2_distilled_fast_profile.py b/examples/inference/basic/basic_ltx2_distilled_fast_profile.py index dae08f5d4..b4e6c42b2 100644 --- a/examples/inference/basic/basic_ltx2_distilled_fast_profile.py +++ b/examples/inference/basic/basic_ltx2_distilled_fast_profile.py @@ -12,15 +12,11 @@ from fastvideo.configs.pipelines.base import PipelineConfig from fastvideo.layers.quantization.nvfp4_config import NVFP4Config from fastvideo.utils import maybe_download_model -VALIDATION_JSON = ( - Path(__file__).resolve().parents[2] / "training" / "finetune" / "ltx2" / "validation.json" -) +VALIDATION_JSON = (Path(__file__).resolve().parents[2] / "training" / "finetune" / "ltx2" / "validation.json") # Override with a local snapshot or converted directory when needed, e.g. # export LTX2_MODEL_PATH=/raid/$USER/hf/FastVideo/LTX2-Distilled-Diffusers -MODEL_ID = os.path.expandvars( - os.path.expanduser(os.getenv("LTX2_MODEL_PATH", "FastVideo/LTX2-Distilled-Diffusers")) -) +MODEL_ID = os.path.expandvars(os.path.expanduser(os.getenv("LTX2_MODEL_PATH", "FastVideo/LTX2-Distilled-Diffusers"))) OUTPUT_DIR = Path("outputs_video/ltx2_distilled_fast_profile") os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" @@ -69,9 +65,7 @@ def print_stage_breakdown( return total -def extract_sr_forward_latency( - result: dict, -) -> tuple[float | None, list[tuple[str, float]], list[str]]: +def extract_sr_forward_latency(result: dict, ) -> tuple[float | None, list[tuple[str, float]], list[str]]: logging_info = result.get("logging_info") if logging_info is None: return None, [], [] @@ -89,12 +83,8 @@ def extract_sr_forward_latency( if sr_match_substr: is_sr_stage = sr_match_substr in stage_name_l else: - is_sr_stage = ( - "srdenoisingstage" in stage_name_l - or "sr_denoising" in stage_name_l - or "upsample" in stage_name_l - or ("refine" in stage_name_l and "denois" in stage_name_l) - ) + is_sr_stage = ("srdenoisingstage" in stage_name_l or "sr_denoising" in stage_name_l + or "upsample" in stage_name_l or ("refine" in stage_name_l and "denois" in stage_name_l)) if not is_sr_stage: continue exec_time = float(stage_metrics.get("execution_time", 0.0)) @@ -161,11 +151,9 @@ def resolve_refine_upsampler_path(model_root: str) -> Path: return candidate checked = "\n".join(f" - {candidate}" for candidate in candidates) - raise FileNotFoundError( - "Could not find an LTX2 refine upsampler directory.\n" - "Checked:\n" - f"{checked}" - ) + raise FileNotFoundError("Could not find an LTX2 refine upsampler directory.\n" + "Checked:\n" + f"{checked}") def main() -> None: @@ -314,19 +302,15 @@ def main() -> None: measured_times = run_times[measured_start_idx:] avg_time = sum(measured_times) / len(measured_times) - print( - f"Average video generation time over {len(measured_times)} runs " - f"(runs {measured_start_idx + 1}-{len(run_times)}, skipping first {warmup_runs} warmup runs): " - f"{avg_time:.2f}s" - ) + print(f"Average video generation time over {len(measured_times)} runs " + f"(runs {measured_start_idx + 1}-{len(run_times)}, skipping first {warmup_runs} warmup runs): " + f"{avg_time:.2f}s") measured_e2e_times = e2e_times[measured_start_idx:] avg_e2e_time = sum(measured_e2e_times) / len(measured_e2e_times) - print( - f"Average end-to-end latency over {len(measured_e2e_times)} runs " - f"(runs {measured_start_idx + 1}-{len(e2e_times)}, skipping first {warmup_runs} warmup runs): " - f"{avg_e2e_time:.2f}s" - ) + print(f"Average end-to-end latency over {len(measured_e2e_times)} runs " + f"(runs {measured_start_idx + 1}-{len(e2e_times)}, skipping first {warmup_runs} warmup runs): " + f"{avg_e2e_time:.2f}s") if sr_forward_times: avg_sr_forward = sum(sr_forward_times) / len(sr_forward_times) @@ -338,10 +322,8 @@ def main() -> None: if non_stage_overhead_times: avg_non_stage_overhead = sum(non_stage_overhead_times) / len(non_stage_overhead_times) - print( - "Average non-stage overhead over " - f"{len(non_stage_overhead_times)} measured runs: {avg_non_stage_overhead:.3f}s" - ) + print("Average non-stage overhead over " + f"{len(non_stage_overhead_times)} measured runs: {avg_non_stage_overhead:.3f}s") else: print("Average non-stage overhead unavailable (no stage timings).") finally: diff --git a/examples/inference/basic/basic_matrixgame2.py b/examples/inference/basic/basic_matrixgame2.py index 70763ecfb..3ffd39011 100644 --- a/examples/inference/basic/basic_matrixgame2.py +++ b/examples/inference/basic/basic_matrixgame2.py @@ -13,24 +13,34 @@ MODEL_VARIANT = "base_distilled_model" # Variant-specific settings VARIANT_CONFIG = { "base_distilled_model": { - "model_path": "FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers", - "keyboard_dim": 4, - "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png", + "model_path": + "FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers", + "keyboard_dim": + 4, + "image_url": + "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png", }, "gta_distilled_model": { - "model_path": "FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers", - "keyboard_dim": 2, - "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/gta_drive/0000.png", + "model_path": + "FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers", + "keyboard_dim": + 2, + "image_url": + "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/gta_drive/0000.png", }, "templerun_distilled_model": { - "model_path": "FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers", - "keyboard_dim": 7, - "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/temple_run/0000.png", + "model_path": + "FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers", + "keyboard_dim": + 7, + "image_url": + "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/temple_run/0000.png", }, } - OUTPUT_PATH = "video_samples_matrixgame2" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. @@ -42,8 +52,8 @@ def main(): config["model_path"], # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory - dit_cpu_offload=True, # DiT need to be offloaded for MoE + use_fsdp_inference=False, # set to True if GPU is out of memory + dit_cpu_offload=True, # DiT need to be offloaded for MoE vae_cpu_offload=False, text_encoder_cpu_offload=True, # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer diff --git a/examples/inference/basic/basic_matrixgame2_streaming.py b/examples/inference/basic/basic_matrixgame2_streaming.py index 0361d5611..454dafb29 100644 --- a/examples/inference/basic/basic_matrixgame2_streaming.py +++ b/examples/inference/basic/basic_matrixgame2_streaming.py @@ -14,27 +14,40 @@ MODEL_VARIANT = "base_distilled_model" # Variant-specific settings VARIANT_CONFIG = { "base_distilled_model": { - "model_path": "FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers", - "keyboard_dim": 4, - "mode": "universal", - "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png", + "model_path": + "FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers", + "keyboard_dim": + 4, + "mode": + "universal", + "image_url": + "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png", }, "gta_distilled_model": { - "model_path": "FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers", - "keyboard_dim": 2, - "mode": "gta_drive", - "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/gta_drive/0000.png", + "model_path": + "FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers", + "keyboard_dim": + 2, + "mode": + "gta_drive", + "image_url": + "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/gta_drive/0000.png", }, "templerun_distilled_model": { - "model_path": "FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers", - "keyboard_dim": 7, - "mode": "templerun", - "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/temple_run/0000.png", + "model_path": + "FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers", + "keyboard_dim": + 7, + "mode": + "templerun", + "image_url": + "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/temple_run/0000.png", }, } - OUTPUT_PATH = "video_samples_matrixgame2" + + async def main(): # FastVideo will automatically use the optimal default arguments for the # model. @@ -46,8 +59,8 @@ async def main(): config["model_path"], # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory - dit_cpu_offload=True, # DiT need to be offloaded for MoE + use_fsdp_inference=False, # set to True if GPU is out of memory + dit_cpu_offload=True, # DiT need to be offloaded for MoE vae_cpu_offload=False, text_encoder_cpu_offload=True, # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer @@ -56,11 +69,8 @@ async def main(): ) max_blocks = 50 - num_frames = 597 - actions = { - "keyboard": torch.zeros((num_frames, config["keyboard_dim"])), - "mouse": torch.zeros((num_frames, 2)) - } + num_frames = 597 + actions = {"keyboard": torch.zeros((num_frames, config["keyboard_dim"])), "mouse": torch.zeros((num_frames, 2))} grid_sizes = torch.tensor([150, 44, 80]) mode = config["mode"] @@ -81,11 +91,11 @@ async def main(): for block_id in range(max_blocks): print(f"\n=== Block {block_id + 1}/{max_blocks} ===") - + action = await get_current_action_async(mode) keyboard_cond, mouse_cond = expand_action_to_frames(action, 12) await generator.step_async(keyboard_cond, mouse_cond) - + if (await asyncio.to_thread(input, "\nContinue? (y/n): ")).lower() == 'n': break diff --git a/examples/inference/basic/basic_minimax_h3_t2v.py b/examples/inference/basic/basic_minimax_h3_t2v.py index c82b11aa5..c3aa16f2c 100644 --- a/examples/inference/basic/basic_minimax_h3_t2v.py +++ b/examples/inference/basic/basic_minimax_h3_t2v.py @@ -30,13 +30,15 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--steps", type=int, default=50) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--num-gpus", type=int, default=4) - parser.add_argument("--torch-compile", action="store_true", - help="torch.compile the DiT transformer path") - parser.add_argument("--compile-mode", default=None, + parser.add_argument("--torch-compile", action="store_true", help="torch.compile the DiT transformer path") + parser.add_argument("--compile-mode", + default=None, help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs') - parser.add_argument("--repeats", type=int, default=1, + parser.add_argument("--repeats", + type=int, + default=1, help="generate N times; with --torch-compile the first run pays " - "compilation, so steady-state is the last repeat") + "compilation, so steady-state is the last repeat") return parser.parse_args() @@ -67,24 +69,24 @@ def main() -> None: )) try: request = GenerationRequest( - prompt=args.prompt, - negative_prompt="", - sampling=SamplingConfig( - height=args.height, - width=args.width, - num_frames=args.num_frames, - fps=24, - num_inference_steps=args.steps, - guidance_scale=1.0, - batch_cfg=False, - seed=args.seed, - ), - output=OutputConfig( - output_path=str(output_dir / "minimax_h3_t2v.mp4"), - save_video=True, - return_frames=False, - ), - ) + prompt=args.prompt, + negative_prompt="", + sampling=SamplingConfig( + height=args.height, + width=args.width, + num_frames=args.num_frames, + fps=24, + num_inference_steps=args.steps, + guidance_scale=1.0, + batch_cfg=False, + seed=args.seed, + ), + output=OutputConfig( + output_path=str(output_dir / "minimax_h3_t2v.mp4"), + save_video=True, + return_frames=False, + ), + ) result = generator.generate(request) print(f"Output written to: {result.video_path}") if result.generation_time is not None: diff --git a/examples/inference/basic/basic_mps.py b/examples/inference/basic/basic_mps.py index 885f945d1..f1d9c214d 100644 --- a/examples/inference/basic/basic_mps.py +++ b/examples/inference/basic/basic_mps.py @@ -1,19 +1,20 @@ from fastvideo import VideoGenerator, PipelineConfig from fastvideo.api.sampling_param import SamplingParam + def main(): config = PipelineConfig.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") config.text_encoder_precisions = ["fp16"] - + generator = VideoGenerator.from_pretrained( "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", pipeline_config=config, - use_fsdp_inference=False, # Disable FSDP for MPS - dit_cpu_offload=True, - text_encoder_cpu_offload=True, - pin_cpu_memory=True, - disable_autocast=False, - num_gpus=1, + use_fsdp_inference=False, # Disable FSDP for MPS + dit_cpu_offload=True, + text_encoder_cpu_offload=True, + pin_cpu_memory=True, + disable_autocast=False, + num_gpus=1, ) # Create sampling parameters with reduced number of frames @@ -23,18 +24,19 @@ def main(): sampling_param.width = 256 prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") - + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") + video = generator.generate_video(prompt, sampling_param=sampling_param) prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") - + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") + video2 = generator.generate_video(prompt2, sampling_param=sampling_param) + if __name__ == "__main__": main() diff --git a/examples/inference/basic/basic_ray.py b/examples/inference/basic/basic_ray.py index 4837e7e01..6d72b430c 100644 --- a/examples/inference/basic/basic_ray.py +++ b/examples/inference/basic/basic_ray.py @@ -3,6 +3,8 @@ from fastvideo import VideoGenerator # from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. @@ -16,27 +18,24 @@ def main(): dit_cpu_offload=False, vae_cpu_offload=False, text_encoder_cpu_offload=True, - pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" + pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" distributed_executor_backend="ray", # image_encoder_cpu_offload=False, ) # Generate videos with the same simple API, regardless of GPU count - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True) # Generate another video with a different prompt, without reloading the # model! - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") + prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True) diff --git a/examples/inference/basic/basic_sd35_t2i.py b/examples/inference/basic/basic_sd35_t2i.py index 6351948c5..0124fc5c2 100644 --- a/examples/inference/basic/basic_sd35_t2i.py +++ b/examples/inference/basic/basic_sd35_t2i.py @@ -7,7 +7,6 @@ import os import re from typing import List - DEFAULT_PROMPTS = [ "a photo of a cat", "a cinematic photo of a red panda wearing a tiny backpack, standing on a rainy neon-lit street at night, shallow depth of field, sharp focus, 35mm, bokeh", @@ -49,7 +48,9 @@ def _remove_existing_outputs(out_dir: str, filename_base: str) -> None: def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description="Run SD3.5 Medium text-to-image with FastVideo VideoGenerator.") - p.add_argument("--model-path", default="stabilityai/stable-diffusion-3.5-medium", help="Path to local diffusers-format SD3.5 weights directory.") + p.add_argument("--model-path", + default="stabilityai/stable-diffusion-3.5-medium", + help="Path to local diffusers-format SD3.5 weights directory.") p.add_argument( "--out-dir", "--outdir", diff --git a/examples/inference/basic/basic_self_forcing_causal.py b/examples/inference/basic/basic_self_forcing_causal.py index 0e6d46e9e..3937226ec 100644 --- a/examples/inference/basic/basic_self_forcing_causal.py +++ b/examples/inference/basic/basic_self_forcing_causal.py @@ -3,6 +3,8 @@ import time from fastvideo import VideoGenerator, SamplingParam OUTPUT_PATH = "video_samples_causal" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. @@ -13,19 +15,18 @@ def main(): model_name, # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory + use_fsdp_inference=False, # set to True if GPU is out of memory text_encoder_cpu_offload=False, dit_cpu_offload=False, ) sampling_param = SamplingParam.from_pretrained(model_name) - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param) + if __name__ == "__main__": main() diff --git a/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py b/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py index 32b60cd86..e75373b54 100644 --- a/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py +++ b/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py @@ -5,6 +5,8 @@ import json # from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_self_forcing_causal_wan2_2_14B_i2v" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. @@ -14,8 +16,8 @@ def main(): "FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers", # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory - dit_cpu_offload=True, # DiT need to be offloaded for MoE + use_fsdp_inference=False, # set to True if GPU is out of memory + dit_cpu_offload=True, # DiT need to be offloaded for MoE dit_precision="fp32", vae_cpu_offload=False, text_encoder_cpu_offload=True, @@ -37,7 +39,11 @@ def main(): for prompt_image_pair in prompt_image_pairs: prompt = prompt_image_pair["prompt"] image_path = prompt_image_pair["image_path"] - _ = generator.generate_video(prompt, image_path=image_path, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param) + _ = generator.generate_video(prompt, + image_path=image_path, + output_path=OUTPUT_PATH, + save_video=True, + sampling_param=sampling_param) if __name__ == "__main__": diff --git a/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py b/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py index 3abdb1b7e..508f84ec2 100644 --- a/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py +++ b/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py @@ -5,6 +5,8 @@ from fastvideo import VideoGenerator # from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_self_forcing_causal_wan2_2_14B_t2v" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. @@ -14,15 +16,17 @@ def main(): "rand0nmr/SFWan2.2-T2V-A14B-Diffusers", # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory - dit_cpu_offload=True, # DiT need to be offloaded for MoE + use_fsdp_inference=False, # set to True if GPU is out of memory + dit_cpu_offload=True, # DiT need to be offloaded for MoE vae_cpu_offload=False, text_encoder_cpu_offload=True, dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125], # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer pin_cpu_memory=True, - init_weights_from_safetensors="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_inference_transformer/", - init_weights_from_safetensors_2="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_2_inference_transformer/", + init_weights_from_safetensors= + "/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_inference_transformer/", + init_weights_from_safetensors_2= + "/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_2_inference_transformer/", num_frame_per_block=7, # image_encoder_cpu_offload=False, ) @@ -31,13 +35,11 @@ def main(): # sampling_param.num_frames = 45 # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" # Generate videos with the same simple API, regardless of GPU count - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") _ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, num_frames=81) if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/inference/basic/basic_stable_audio_inpaint.py b/examples/inference/basic/basic_stable_audio_inpaint.py index 126bbc54a..0424edbb1 100644 --- a/examples/inference/basic/basic_stable_audio_inpaint.py +++ b/examples/inference/basic/basic_stable_audio_inpaint.py @@ -54,16 +54,15 @@ PROMPT = "Steady lo-fi hip hop drum loop with vinyl crackle." # ...) you want to extend or repair. The pipeline raises if a mask is # passed without a reference, so this must be a real path. REFERENCE_AUDIO_PATH = "path/to/your/loop.wav" -KEEP_SECONDS = 6.0 # first KEEP_SECONDS preserved exactly -TOTAL_SECONDS = 12.0 # extend the loop to this duration +KEEP_SECONDS = 6.0 # first KEEP_SECONDS preserved exactly +TOTAL_SECONDS = 12.0 # extend the loop to this duration def main() -> None: if not os.path.isfile(REFERENCE_AUDIO_PATH): - raise FileNotFoundError( - f"REFERENCE_AUDIO_PATH={REFERENCE_AUDIO_PATH!r} does not exist. " - "Edit this script to point at a real audio file (wav/mp3/mp4/" - "m4a/flac) before running.") + raise FileNotFoundError(f"REFERENCE_AUDIO_PATH={REFERENCE_AUDIO_PATH!r} does not exist. " + "Edit this script to point at a real audio file (wav/mp3/mp4/" + "m4a/flac) before running.") generator = VideoGenerator.from_pretrained( "FastVideo/stable-audio-open-1.0-Diffusers", num_gpus=1, diff --git a/examples/inference/basic/basic_turbodiffusion.py b/examples/inference/basic/basic_turbodiffusion.py index b3e314869..6e85a734e 100644 --- a/examples/inference/basic/basic_turbodiffusion.py +++ b/examples/inference/basic/basic_turbodiffusion.py @@ -15,19 +15,17 @@ def main() -> None: "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers", # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory + use_fsdp_inference=False, # set to True if GPU is out of memory - # set to false if using RTX 4090 + # set to false if using RTX 4090 # pin_cpu_memory=False, ) # Generate videos with the same simple API, regardless of GPU count # TurboDiffusion defaults: guidance_scale=1.0 and num_inference_steps=4 (from config) - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") video = generator.generate_video( prompt, output_path=OUTPUT_PATH, @@ -36,13 +34,11 @@ def main() -> None: ) # Generate another video with a different prompt, without reloading the model! - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic." - ) + prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") video2 = generator.generate_video( prompt2, output_path=OUTPUT_PATH, diff --git a/examples/inference/basic/basic_turbodiffusion_14b.py b/examples/inference/basic/basic_turbodiffusion_14b.py index 50cfd337b..ebaa39c3b 100644 --- a/examples/inference/basic/basic_turbodiffusion_14b.py +++ b/examples/inference/basic/basic_turbodiffusion_14b.py @@ -17,11 +17,9 @@ def main() -> None: num_gpus=2, ) - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") video = generator.generate_video( prompt, output_path=OUTPUT_PATH, @@ -30,13 +28,11 @@ def main() -> None: ) # Generate another video with a different prompt, without reloading the model! - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic." - ) + prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") video2 = generator.generate_video( prompt2, output_path=OUTPUT_PATH, diff --git a/examples/inference/basic/basic_turbodiffusion_i2v.py b/examples/inference/basic/basic_turbodiffusion_i2v.py index 655dbbbdd..b28092f21 100644 --- a/examples/inference/basic/basic_turbodiffusion_i2v.py +++ b/examples/inference/basic/basic_turbodiffusion_i2v.py @@ -18,12 +18,13 @@ def main() -> None: ) # Example prompt and image for I2V - prompt = ("Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside.") + prompt = ( + "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." + ) # Use an example image path image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG" - video = generator.generate_video( prompt, image_path=image_path, diff --git a/examples/inference/basic/basic_wan2_2.py b/examples/inference/basic/basic_wan2_2.py index e6758347c..02d8d8d38 100644 --- a/examples/inference/basic/basic_wan2_2.py +++ b/examples/inference/basic/basic_wan2_2.py @@ -3,6 +3,8 @@ from fastvideo import VideoGenerator # from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_wan2_2_14B_t2v" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. @@ -12,8 +14,8 @@ def main(): "Wan-AI/Wan2.2-T2V-A14B-Diffusers", # FastVideo will automatically handle distributed setup num_gpus=2, - use_fsdp_inference=False, # set to True if GPU is out of memory - dit_cpu_offload=True, # DiT need to be offloaded for MoE + use_fsdp_inference=False, # set to True if GPU is out of memory + dit_cpu_offload=True, # DiT need to be offloaded for MoE vae_cpu_offload=False, text_encoder_cpu_offload=True, # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer @@ -25,24 +27,31 @@ def main(): # sampling_param.num_frames = 45 # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" # Generate videos with the same simple API, regardless of GPU count - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) - _ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, height=720, width=1280, num_frames=81) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") + _ = generator.generate_video(prompt, + output_path=OUTPUT_PATH, + save_video=True, + height=720, + width=1280, + num_frames=81) # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/") # Generate another video with a different prompt, without reloading the # model! - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") - _ = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, height=720, width=1280, num_frames=81) + prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") + _ = generator.generate_video(prompt2, + output_path=OUTPUT_PATH, + save_video=True, + height=720, + width=1280, + num_frames=81) if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/inference/basic/basic_wan2_2_Fun.py b/examples/inference/basic/basic_wan2_2_Fun.py index b3d6f1229..c41699a46 100644 --- a/examples/inference/basic/basic_wan2_2_Fun.py +++ b/examples/inference/basic/basic_wan2_2_Fun.py @@ -4,6 +4,8 @@ from fastvideo import VideoGenerator OUTPUT_PATH = "video_samples_wan2_1_Fun" OUTPUT_NAME = "wan2.1_test" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. @@ -14,8 +16,8 @@ def main(): # "alibaba-pai/Wan2.2-Fun-A14B-Control", # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory - dit_cpu_offload=True, # DiT need to be offloaded for MoE + use_fsdp_inference=False, # set to True if GPU is out of memory + dit_cpu_offload=True, # DiT need to be offloaded for MoE vae_cpu_offload=False, text_encoder_cpu_offload=True, # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer @@ -30,7 +32,14 @@ def main(): image_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/8.png" control_video_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/pose.mp4" - video = generator.generate_video(prompt, negative_prompt=negative_prompt, image_path=image_path, video_path=control_video_path, output_path=OUTPUT_PATH, output_video_name=OUTPUT_NAME, save_video=True) + video = generator.generate_video(prompt, + negative_prompt=negative_prompt, + image_path=image_path, + video_path=control_video_path, + output_path=OUTPUT_PATH, + output_video_name=OUTPUT_NAME, + save_video=True) + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/inference/basic/basic_wan2_2_i2v.py b/examples/inference/basic/basic_wan2_2_i2v.py index 15babf2aa..a67d57b84 100644 --- a/examples/inference/basic/basic_wan2_2_i2v.py +++ b/examples/inference/basic/basic_wan2_2_i2v.py @@ -3,6 +3,8 @@ from fastvideo import VideoGenerator # from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_wan2_2_14B_i2v" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. @@ -12,8 +14,8 @@ def main(): "Wan-AI/Wan2.2-I2V-A14B-Diffusers", # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory - dit_cpu_offload=True, # DiT need to be offloaded for MoE + use_fsdp_inference=False, # set to True if GPU is out of memory + dit_cpu_offload=True, # DiT need to be offloaded for MoE vae_cpu_offload=False, text_encoder_cpu_offload=True, # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer @@ -24,7 +26,14 @@ def main(): prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG" - video = generator.generate_video(prompt, image_path=image_path, output_path=OUTPUT_PATH, save_video=True, height=832, width=480, num_frames=81) + video = generator.generate_video(prompt, + image_path=image_path, + output_path=OUTPUT_PATH, + save_video=True, + height=832, + width=480, + num_frames=81) + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/inference/basic/basic_wan2_2_ti2v.py b/examples/inference/basic/basic_wan2_2_ti2v.py index df0c18848..c9ca208e4 100644 --- a/examples/inference/basic/basic_wan2_2_ti2v.py +++ b/examples/inference/basic/basic_wan2_2_ti2v.py @@ -1,6 +1,8 @@ from fastvideo import VideoGenerator OUTPUT_PATH = "video_samples_wan2_2_5B_ti2v" + + def main(): # FastVideo will automatically use the optimal default arguments for the # model. @@ -11,11 +13,11 @@ def main(): model_name, # FastVideo will automatically handle distributed setup num_gpus=1, - use_fsdp_inference=False, # set to True if GPU is out of memory + use_fsdp_inference=False, # set to True if GPU is out of memory dit_cpu_offload=True, vae_cpu_offload=False, text_encoder_cpu_offload=True, - pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" + pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" # image_encoder_cpu_offload=False, ) @@ -28,14 +30,13 @@ def main(): # model! # T2V mode - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") + prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic.") video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True) if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/inference/basic/basic_zimage.py b/examples/inference/basic/basic_zimage.py index 0ecbf5d87..732894163 100644 --- a/examples/inference/basic/basic_zimage.py +++ b/examples/inference/basic/basic_zimage.py @@ -20,13 +20,11 @@ from fastvideo.api import ( SamplingConfig, ) - DEFAULT_PROMPT = ( "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. " "Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. " "Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, " - "silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights." -) + "silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights.") DEFAULT_REVISION = "f332072aa78be7aecdf3ee76d5c247082da564a6" diff --git a/examples/inference/eval/all_metrics_demo.py b/examples/inference/eval/all_metrics_demo.py index 471ff489a..0b6eda9c2 100644 --- a/examples/inference/eval/all_metrics_demo.py +++ b/examples/inference/eval/all_metrics_demo.py @@ -35,13 +35,11 @@ import torch # Generation — make one LTX2 video to evaluate. # --------------------------------------------------------------------- -PROMPT = ( - "A warm sunny backyard. The camera starts in a tight cinematic close-up " - "of a woman and a man in their 30s, facing each other with serious " - "expressions. The woman, emotional and dramatic, says softly, \"That's " - "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " - "annoyed: \"Stop being so dramatic, Jess.\"" -) +PROMPT = ("A warm sunny backyard. The camera starts in a tight cinematic close-up " + "of a woman and a man in their 30s, facing each other with serious " + "expressions. The woman, emotional and dramatic, says softly, \"That's " + "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " + "annoyed: \"Stop being so dramatic, Jess.\"") OUTPUT_PATH = "fastvideo/tests/eval/asset/ltx2.mp4" N_DUP = 4 # how many times to duplicate the video for the gen/ref corpora @@ -77,6 +75,7 @@ def generate_one_ltx2_video() -> str: # Eval — the point of the script. 4 lines from "two paths" to results. # --------------------------------------------------------------------- + def _all_registered_metrics() -> list[str]: """Every metric in the registry, sorted. Combined with ``skip_missing_deps=True`` this is the "run everything that works in @@ -103,8 +102,7 @@ def score_all_metrics(video_path: str) -> None: torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() t_init0 = time.perf_counter() - ev = create_evaluator(metrics=_all_registered_metrics(), - device="cuda:0", num_gpus=1, skip_missing_deps=True) + ev = create_evaluator(metrics=_all_registered_metrics(), device="cuda:0", num_gpus=1, skip_missing_deps=True) t_init1 = time.perf_counter() samples = samples_from(video=gen_dir, reference=ref_dir, text_prompt=PROMPT, fps=24.0, extract_audio=True) # auto-extract audio track from videos diff --git a/examples/inference/eval/basic_ltx2_audio_eval.py b/examples/inference/eval/basic_ltx2_audio_eval.py index 986ebea22..3bf314c07 100644 --- a/examples/inference/eval/basic_ltx2_audio_eval.py +++ b/examples/inference/eval/basic_ltx2_audio_eval.py @@ -23,14 +23,12 @@ import torch from fastvideo import VideoGenerator from fastvideo.eval import create_evaluator -PROMPT = ( - "A warm sunny backyard. The camera starts in a tight cinematic close-up " - "of a woman and a man in their 30s, facing each other with serious " - "expressions. The woman, emotional and dramatic, says softly, \"That's " - "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " - "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " - "then mutters defensively, \"He's just having fun.\"" -) +PROMPT = ("A warm sunny backyard. The camera starts in a tight cinematic close-up " + "of a woman and a man in their 30s, facing each other with serious " + "expressions. The woman, emotional and dramatic, says softly, \"That's " + "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " + "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " + "then mutters defensively, \"He's just having fun.\"") METRICS = [ "audio.clap_score", diff --git a/examples/inference/eval/basic_ltx2_eval.py b/examples/inference/eval/basic_ltx2_eval.py index caeb2293c..55c2b6503 100644 --- a/examples/inference/eval/basic_ltx2_eval.py +++ b/examples/inference/eval/basic_ltx2_eval.py @@ -25,19 +25,17 @@ from fastvideo import VideoGenerator from fastvideo.eval import Evaluator from fastvideo.eval.io import build_eval_kwargs -PROMPT = ( - "A warm sunny backyard. The camera starts in a tight cinematic close-up " - "of a woman and a man in their 30s, facing each other with serious " - "expressions. The woman, emotional and dramatic, says softly, \"That's " - "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " - "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " - "then mutters defensively, \"He's just having fun.\" The camera slowly " - "pans right, revealing the grandfather in the garden wearing enormous " - "butterfly wings, waving his arms in the air like he's trying to take " - "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. " - "The woman covers her face, on the verge of tears. The tone is deadpan, " - "absurd, and quietly tragic." -) +PROMPT = ("A warm sunny backyard. The camera starts in a tight cinematic close-up " + "of a woman and a man in their 30s, facing each other with serious " + "expressions. The woman, emotional and dramatic, says softly, \"That's " + "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " + "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " + "then mutters defensively, \"He's just having fun.\" The camera slowly " + "pans right, revealing the grandfather in the garden wearing enormous " + "butterfly wings, waving his arms in the air like he's trying to take " + "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. " + "The woman covers her face, on the verge of tears. The tone is deadpan, " + "absurd, and quietly tragic.") # VBench sub-metrics meaningful for an arbitrary text→video sample # (just the generated frames, optionally fps + the source prompt). @@ -45,14 +43,14 @@ PROMPT = ( # vbench.scene, ...) are excluded — they need prompts built to a # specific schema. METRICS = [ - "vbench.aesthetic_quality", # CLIP + LAION aesthetic head - "vbench.subject_consistency", # DINO frame-to-first cosine + "vbench.aesthetic_quality", # CLIP + LAION aesthetic head + "vbench.subject_consistency", # DINO frame-to-first cosine "vbench.background_consistency", # DINO on background patches - "vbench.imaging_quality", # pyiqa MUSIQ - "vbench.temporal_flickering", # pixel-wise frame deltas - "vbench.motion_smoothness", # AMT frame interpolator residual - "vbench.dynamic_degree", # RAFT optical-flow magnitude (needs fps) - "vbench.overall_consistency", # ViCLIP video↔prompt similarity + "vbench.imaging_quality", # pyiqa MUSIQ + "vbench.temporal_flickering", # pixel-wise frame deltas + "vbench.motion_smoothness", # AMT frame interpolator residual + "vbench.dynamic_degree", # RAFT optical-flow magnitude (needs fps) + "vbench.overall_consistency", # ViCLIP video↔prompt similarity ] diff --git a/examples/inference/eval/bench_physics_iq.py b/examples/inference/eval/bench_physics_iq.py index a33298e67..1feb2e771 100644 --- a/examples/inference/eval/bench_physics_iq.py +++ b/examples/inference/eval/bench_physics_iq.py @@ -41,9 +41,8 @@ def _expected_filename(row: dict) -> str: return row["auxiliary_info"]["expected_gen_filename"] -def _generate_videos(rows: list[dict], videos_dir: Path, - model: str, num_gpus: int, - num_frames: int, height: int, width: int) -> None: +def _generate_videos(rows: list[dict], videos_dir: Path, model: str, num_gpus: int, num_frames: int, height: int, + width: int) -> None: from fastvideo import VideoGenerator videos_dir.mkdir(parents=True, exist_ok=True) @@ -59,36 +58,43 @@ def _generate_videos(rows: list[dict], videos_dir: Path, try: for row, out_path in todo: gen.generate_video( - prompt=row["prompt"], output_path=str(out_path), save_video=True, - num_frames=num_frames, height=height, width=width, + prompt=row["prompt"], + output_path=str(out_path), + save_video=True, + num_frames=num_frames, + height=height, + width=width, ) finally: gen.shutdown() def main() -> None: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--dataset-root", type=Path, default=None, + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dataset-root", + type=Path, + default=None, help="Path to a pre-downloaded Physics-IQ release. " - "Defaults to ${FASTVIDEO_EVAL_CACHE}/datasets/physics_iq, " - "auto-fetching missing assets from the public bucket.") - p.add_argument("--videos-dir", type=Path, + "Defaults to ${FASTVIDEO_EVAL_CACHE}/datasets/physics_iq, " + "auto-fetching missing assets from the public bucket.") + p.add_argument("--videos-dir", + type=Path, default=Path("outputs_video/bench_physics_iq"), help="Where to read/write generated videos.") - p.add_argument("--limit", type=int, default=None, - help="Truncate to first N scenarios for smoke runs.") + p.add_argument("--limit", type=int, default=None, help="Truncate to first N scenarios for smoke runs.") p.add_argument("--num-gpus", type=int, default=1) - p.add_argument("--model", default="Davids048/LTX2-Base-Diffusers", + p.add_argument("--model", + default="Davids048/LTX2-Base-Diffusers", help="HF repo id of the text→video generator to use.") p.add_argument("--num-frames", type=int, default=121) p.add_argument("--height", type=int, default=1088) p.add_argument("--width", type=int, default=1920) - p.add_argument("--skip-generation", action="store_true", - help="Re-score existing videos under --videos-dir.") - p.add_argument("--scores-out", type=Path, default=None, + p.add_argument("--skip-generation", action="store_true", help="Re-score existing videos under --videos-dir.") + p.add_argument("--scores-out", + type=Path, + default=None, help="Where to write per-scenario scores (JSON). " - "Defaults to /scores.json.") + "Defaults to /scores.json.") args = p.parse_args() # 1. Walk the Physics-IQ corpus. Pass --limit to the dataset @@ -100,8 +106,13 @@ def main() -> None: # 2. Generate (or reuse) one mp4 per scenario. if not args.skip_generation: _generate_videos( - rows, args.videos_dir, args.model, args.num_gpus, - args.num_frames, args.height, args.width, + rows, + args.videos_dir, + args.model, + args.num_gpus, + args.num_frames, + args.height, + args.width, ) # 3. Score each scenario. The metric reads file paths directly out @@ -126,28 +137,26 @@ def main() -> None: # 4. Aggregate per the upstream scoring recipe. metric = get_metric("physics_iq") - components = metric.aggregate_components( - [r["physics_iq"] for r in all_results] - ) + components = metric.aggregate_components([r["physics_iq"] for r in all_results]) print() print("=== Physics-IQ aggregate ===") for name, value in components.items(): print(f" {name:24s} {value:.4f}") - detailed = [ - { - "scenario": row["auxiliary_info"]["scenario_id"], - "view": row["view"], - "scenario_name": row["auxiliary_info"]["scenario_name"], - "score": results["physics_iq"].score, - } - for row, results in zip(matched, all_results) - ] + detailed = [{ + "scenario": row["auxiliary_info"]["scenario_id"], + "view": row["view"], + "scenario_name": row["auxiliary_info"]["scenario_name"], + "score": results["physics_iq"].score, + } for row, results in zip(matched, all_results)] out = args.scores_out or (args.videos_dir / "scores.json") out.parent.mkdir(parents=True, exist_ok=True) out.write_text(json.dumps( - {"aggregate": components, "per_scenario": detailed}, + { + "aggregate": components, + "per_scenario": detailed + }, indent=2, )) print(f"\n[done] per-scenario scores → {out}") diff --git a/examples/inference/eval/bench_vbench.py b/examples/inference/eval/bench_vbench.py index d2a5d8e43..777d1c346 100644 --- a/examples/inference/eval/bench_vbench.py +++ b/examples/inference/eval/bench_vbench.py @@ -39,9 +39,8 @@ def _slugify(prompt: str, max_len: int = 100) -> str: return re.sub(r"\s+", " ", s) or "output" -def _generate_videos(prompts: list[str], videos_dir: Path, - model: str, num_gpus: int, - num_frames: int, height: int, width: int) -> None: +def _generate_videos(prompts: list[str], videos_dir: Path, model: str, num_gpus: int, num_frames: int, height: int, + width: int) -> None: from fastvideo import VideoGenerator videos_dir.mkdir(parents=True, exist_ok=True) @@ -57,53 +56,60 @@ def _generate_videos(prompts: list[str], videos_dir: Path, try: for prompt, out_path in todo: gen.generate_video( - prompt=prompt, output_path=str(out_path), save_video=True, - num_frames=num_frames, height=height, width=width, + prompt=prompt, + output_path=str(out_path), + save_video=True, + num_frames=num_frames, + height=height, + width=width, ) finally: gen.shutdown() def main() -> None: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--dimensions", default="aesthetic_quality,subject_consistency", + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dimensions", + default="aesthetic_quality,subject_consistency", help="Comma-separated VBench dimensions (or 'all').") - p.add_argument("--limit", type=int, default=None, - help="Truncate to first N prompts for smoke runs.") - p.add_argument("--videos-dir", type=Path, - default=Path("outputs_video/bench_vbench")) + p.add_argument("--limit", type=int, default=None, help="Truncate to first N prompts for smoke runs.") + p.add_argument("--videos-dir", type=Path, default=Path("outputs_video/bench_vbench")) p.add_argument("--num-gpus", type=int, default=1) - p.add_argument("--model", default="Davids048/LTX2-Base-Diffusers", + p.add_argument("--model", + default="Davids048/LTX2-Base-Diffusers", help="HF repo id of the text→video generator to use.") p.add_argument("--num-frames", type=int, default=121) p.add_argument("--height", type=int, default=1088) p.add_argument("--width", type=int, default=1920) - p.add_argument("--fps", type=float, default=24.0, - help="Frame-rate annotation passed to fps-aware metrics.") - p.add_argument("--skip-generation", action="store_true", + p.add_argument("--fps", type=float, default=24.0, help="Frame-rate annotation passed to fps-aware metrics.") + p.add_argument("--skip-generation", + action="store_true", help="Re-score existing videos under --videos-dir without " - "regenerating.") - p.add_argument("--scores-out", type=Path, default=None, + "regenerating.") + p.add_argument("--scores-out", + type=Path, + default=None, help="Where to dump per-prompt scores as JSON. " - "Defaults to /scores.json.") + "Defaults to /scores.json.") args = p.parse_args() # 1. Pull prompts from VBench. - dims_arg: list[str] | str = ( - args.dimensions if args.dimensions == "all" - else [d.strip() for d in args.dimensions.split(",") if d.strip()] - ) + dims_arg: list[str] | str = (args.dimensions if args.dimensions == "all" else + [d.strip() for d in args.dimensions.split(",") if d.strip()]) ds = get_dataset("vbench", dimensions=dims_arg) - rows = list(ds)[: args.limit] + rows = list(ds)[:args.limit] print(f"[load] VBench: {len(rows)} prompts across {ds.dimensions}") # 2. Generate (or reuse) one mp4 per prompt. if not args.skip_generation: _generate_videos( [row["prompt"] for row in rows], - args.videos_dir, args.model, args.num_gpus, - args.num_frames, args.height, args.width, + args.videos_dir, + args.model, + args.num_gpus, + args.num_frames, + args.height, + args.width, ) # 3. Score each video against the requested vbench sub-metrics. @@ -122,7 +128,7 @@ def main() -> None: samples.append({ "video": str(video_path), "fps": args.fps, - **row, # prompt / aux / dims + **row, # prompt / aux / dims }) matched_rows.append(row) diff --git a/examples/inference/eval/eval_fvd.py b/examples/inference/eval/eval_fvd.py index 77e7c17c2..f66c9da8a 100644 --- a/examples/inference/eval/eval_fvd.py +++ b/examples/inference/eval/eval_fvd.py @@ -35,20 +35,29 @@ from fastvideo.eval import create_evaluator, samples_from def main() -> None: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--gen-dir", type=Path, required=True, + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--gen-dir", + type=Path, + required=True, help="Directory of generated videos (.mp4, .avi, .mov, .mkv, .gif).") - p.add_argument("--reference-dir", type=Path, default=None, + p.add_argument("--reference-dir", + type=Path, + default=None, help="Directory of reference videos. Omit to score against the cached " - "reference features (built on a previous run).") + "reference features (built on a previous run).") p.add_argument("--device", default="cuda:0" if torch.cuda.is_available() else "cpu") - p.add_argument("--num-gpus", type=int, default=1, + p.add_argument("--num-gpus", + type=int, + default=1, help="Number of GPU replicas. >1 fans extraction out across devices.") - p.add_argument("--cache-path", type=Path, default=None, + p.add_argument("--cache-path", + type=Path, + default=None, help="Override the reference-feature cache path. " - "Defaults to ${FASTVIDEO_EVAL_CACHE}/fvd/real_features_i3d.pt.") - p.add_argument("--output", type=Path, default=None, + "Defaults to ${FASTVIDEO_EVAL_CACHE}/fvd/real_features_i3d.pt.") + p.add_argument("--output", + type=Path, + default=None, help="Write the result as JSON to this path (default: stdout only).") args = p.parse_args() diff --git a/examples/inference/eval/eval_ltx2_vbench.py b/examples/inference/eval/eval_ltx2_vbench.py index b2d7760e6..8d72e04a2 100644 --- a/examples/inference/eval/eval_ltx2_vbench.py +++ b/examples/inference/eval/eval_ltx2_vbench.py @@ -36,57 +36,55 @@ from fastvideo import VideoGenerator from fastvideo.eval import create_evaluator from fastvideo.eval.io import load_video - -PROMPT = ( - "A warm sunny backyard. The camera starts in a tight cinematic close-up " - "of a woman and a man in their 30s, facing each other with serious " - "expressions. The woman, emotional and dramatic, says softly, \"That's " - "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " - "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " - "then mutters defensively, \"He's just having fun.\" The camera slowly " - "pans right, revealing the grandfather in the garden wearing enormous " - "butterfly wings, waving his arms in the air like he's trying to take " - "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. " - "The woman covers her face, on the verge of tears. The tone is deadpan, " - "absurd, and quietly tragic." -) +PROMPT = ("A warm sunny backyard. The camera starts in a tight cinematic close-up " + "of a woman and a man in their 30s, facing each other with serious " + "expressions. The woman, emotional and dramatic, says softly, \"That's " + "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " + "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " + "then mutters defensively, \"He's just having fun.\" The camera slowly " + "pans right, revealing the grandfather in the garden wearing enormous " + "butterfly wings, waving his arms in the air like he's trying to take " + "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. " + "The woman covers her face, on the verge of tears. The tone is deadpan, " + "absurd, and quietly tragic.") DEFAULT_METRICS = [ # No-input metrics: just need the generated frames. - "vbench.aesthetic_quality", # CLIP + LAION aesthetic head - "vbench.subject_consistency", # DINO frame-to-first cosine + "vbench.aesthetic_quality", # CLIP + LAION aesthetic head + "vbench.subject_consistency", # DINO frame-to-first cosine "vbench.background_consistency", # DINO on background patches - "vbench.imaging_quality", # pyiqa MUSIQ - "vbench.temporal_flickering", # pixel-wise frame deltas - "vbench.motion_smoothness", # AMT frame interpolator residual + "vbench.imaging_quality", # pyiqa MUSIQ + "vbench.temporal_flickering", # pixel-wise frame deltas + "vbench.motion_smoothness", # AMT frame interpolator residual # Need fps annotation: - "vbench.dynamic_degree", # RAFT optical-flow magnitude + "vbench.dynamic_degree", # RAFT optical-flow magnitude # Need the source prompt: - "vbench.overall_consistency", # ViCLIP video↔prompt similarity + "vbench.overall_consistency", # ViCLIP video↔prompt similarity ] def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--model", default="Davids048/LTX2-Base-Diffusers", - help="HF repo id of the LTX2 checkpoint.") + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--model", default="Davids048/LTX2-Base-Diffusers", help="HF repo id of the LTX2 checkpoint.") p.add_argument("--num-gpus", type=int, default=1) - p.add_argument("--output", default="outputs_video/ltx2_eval/clip.mp4", - help="Where to save the generated mp4.") + p.add_argument("--output", default="outputs_video/ltx2_eval/clip.mp4", help="Where to save the generated mp4.") p.add_argument("--num-frames", type=int, default=121) p.add_argument("--height", type=int, default=1088) p.add_argument("--width", type=int, default=1920) p.add_argument("--prompt", default=PROMPT) - p.add_argument("--fps", type=float, default=24.0, + p.add_argument("--fps", + type=float, + default=24.0, help="Frame-rate annotation passed to fps-aware metrics " - "(e.g. vbench.dynamic_degree). LTX2 outputs at 24 fps " - "by default.") - p.add_argument("--metrics", default=",".join(DEFAULT_METRICS), + "(e.g. vbench.dynamic_degree). LTX2 outputs at 24 fps " + "by default.") + p.add_argument("--metrics", + default=",".join(DEFAULT_METRICS), help="Comma-separated metric names. Pass 'all' for every " - "registered metric, or e.g. 'vbench' for the whole group.") + "registered metric, or e.g. 'vbench' for the whole group.") p.add_argument("--scores-out", default="outputs_video/ltx2_eval/scores.json") - p.add_argument("--skip-generation", action="store_true", + p.add_argument("--skip-generation", + action="store_true", help="Reuse an existing --output video instead of regenerating.") return p.parse_args() @@ -115,11 +113,10 @@ def generate(args: argparse.Namespace) -> Path: return out -def evaluate_video(video_path: Path, prompt: str, fps: float, - metric_names) -> dict: +def evaluate_video(video_path: Path, prompt: str, fps: float, metric_names) -> dict: print(f"[eval] loading video from {video_path}...") - video = load_video(str(video_path)) # (T, C, H, W) in [0, 1] - video = video.unsqueeze(0) # → (1, T, C, H, W) + video = load_video(str(video_path)) # (T, C, H, W) in [0, 1] + video = video.unsqueeze(0) # → (1, T, C, H, W) print(f"[eval] building evaluator: {metric_names}") evaluator = create_evaluator(metrics=metric_names, device="cuda") @@ -132,12 +129,9 @@ def evaluate_video(video_path: Path, prompt: str, fps: float, ) if isinstance(results, list): - results = results[0] # batch of 1 + results = results[0] # batch of 1 - return { - name: {"score": r.score, "details": r.details} - for name, r in results.items() - } + return {name: {"score": r.score, "details": r.details} for name, r in results.items()} def main() -> None: @@ -157,7 +151,11 @@ def main() -> None: out = Path(args.scores_out) out.parent.mkdir(parents=True, exist_ok=True) out.write_text(json.dumps( - {"video": str(video_path), "prompt": args.prompt, "scores": scores}, + { + "video": str(video_path), + "prompt": args.prompt, + "scores": scores + }, indent=2, )) print(f"[done] scores written to {out}") diff --git a/examples/inference/eval/eval_third_person_separation.py b/examples/inference/eval/eval_third_person_separation.py index cc0029d62..2549349d3 100644 --- a/examples/inference/eval/eval_third_person_separation.py +++ b/examples/inference/eval/eval_third_person_separation.py @@ -23,19 +23,25 @@ def _by_stem(directory: Path, exts: set[str]) -> dict[str, Path]: def main() -> None: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--candidate-dir", type=Path, default=None, - help="Directory of candidate clips (directory mode).") - p.add_argument("--reference", action="append", default=[], metavar="NAME=DIR", + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--candidate-dir", type=Path, default=None, help="Directory of candidate clips (directory mode).") + p.add_argument("--reference", + action="append", + default=[], + metavar="NAME=DIR", help="Baseline directory, repeatable: 'name=dir' or bare 'dir'.") - p.add_argument("--image-dir", type=Path, default=None, + p.add_argument("--image-dir", + type=Path, + default=None, help="Optional first-frame images, matched to clips by stem.") - p.add_argument("--prompts-json", type=Path, default=None, - help="Optional {stem: control-signal text} JSON.") - p.add_argument("--actions-json", type=Path, default=None, + p.add_argument("--prompts-json", type=Path, default=None, help="Optional {stem: control-signal text} JSON.") + p.add_argument("--actions-json", + type=Path, + default=None, help="Optional {stem: action-label} JSON for the per-action breakdown.") - p.add_argument("--manifest", type=Path, default=None, + p.add_argument("--manifest", + type=Path, + default=None, help="JSON list of {baseline, video_path, reference_path, ...} rows.") p.add_argument("--output", type=Path, default=None) args = p.parse_args() @@ -44,8 +50,7 @@ def main() -> None: by_baseline: dict[str, list[dict]] = defaultdict(list) if args.manifest is not None: for row in json.loads(args.manifest.read_text()): - by_baseline[row.get("baseline", "baseline")].append( - {k: v for k, v in row.items() if k != "baseline"}) + by_baseline[row.get("baseline", "baseline")].append({k: v for k, v in row.items() if k != "baseline"}) elif args.candidate_dir is not None and args.reference: cands = _by_stem(args.candidate_dir, VIDEO_EXTS) images = _by_stem(args.image_dir, IMAGE_EXTS) if args.image_dir else {} diff --git a/examples/inference/eval/score_folder.py b/examples/inference/eval/score_folder.py index 0217ab801..3853e9ce5 100644 --- a/examples/inference/eval/score_folder.py +++ b/examples/inference/eval/score_folder.py @@ -34,16 +34,15 @@ def _list_videos(directory: Path) -> list[Path]: def main() -> None: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--videos", type=Path, required=True, - help="Directory of generated videos.") - p.add_argument("--reference-dir", type=Path, default=None, + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--videos", type=Path, required=True, help="Directory of generated videos.") + p.add_argument("--reference-dir", + type=Path, + default=None, help="Directory of reference videos with matching stems.") p.add_argument("--metrics", default="vbench.aesthetic_quality") p.add_argument("--num-gpus", type=int, default=1) - p.add_argument("--fps", type=float, default=None, - help="Frame-rate annotation for fps-aware metrics.") + p.add_argument("--fps", type=float, default=None, help="Frame-rate annotation for fps-aware metrics.") p.add_argument("--output", type=Path, default=Path("scores.json")) args = p.parse_args() @@ -52,10 +51,8 @@ def main() -> None: raise SystemExit(f"No videos under {args.videos}") print(f"Found {len(video_paths)} videos in {args.videos}") - metrics: list[str] | str = ( - args.metrics if args.metrics == "all" - else [m.strip() for m in args.metrics.split(",") if m.strip()] - ) + metrics: list[str] | str = (args.metrics + if args.metrics == "all" else [m.strip() for m in args.metrics.split(",") if m.strip()]) evaluator = create_evaluator(metrics=metrics, num_gpus=args.num_gpus) # Build per-video sample dicts holding *paths*, not pre-loaded @@ -79,13 +76,13 @@ def main() -> None: all_results = evaluator.evaluate(samples=samples) evaluator.shutdown() - payload = [ - { - "video": str(vp), - "scores": {name: r.score for name, r in results.items()}, - } - for vp, results in zip(video_paths, all_results) - ] + payload = [{ + "video": str(vp), + "scores": { + name: r.score + for name, r in results.items() + }, + } for vp, results in zip(video_paths, all_results)] args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(payload, indent=2)) print(f"Wrote {args.output}") diff --git a/examples/inference/eval/score_video.py b/examples/inference/eval/score_video.py index c9e12dc4d..757a01d2c 100644 --- a/examples/inference/eval/score_video.py +++ b/examples/inference/eval/score_video.py @@ -25,26 +25,26 @@ from fastvideo.eval.io import load_video def main() -> None: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument("--video", required=True, help="Path to the generated mp4.") - p.add_argument("--reference", default=None, - help="Optional path to a reference mp4 (for paired metrics).") - p.add_argument("--metrics", default="common.psnr,common.ssim", + p.add_argument("--reference", default=None, help="Optional path to a reference mp4 (for paired metrics).") + p.add_argument("--metrics", + default="common.psnr,common.ssim", help="Comma-separated metric names, or a group name like 'vbench'.") p.add_argument("--device", default="cuda:0") - p.add_argument("--text-prompt", default=None, + p.add_argument("--text-prompt", + default=None, help="Text prompt for prompt-aware metrics " - "(vbench.overall_consistency, etc.).") - p.add_argument("--fps", type=float, default=None, + "(vbench.overall_consistency, etc.).") + p.add_argument("--fps", + type=float, + default=None, help="Frame-rate annotation for fps-aware metrics " - "(vbench.dynamic_degree, etc.).") + "(vbench.dynamic_degree, etc.).") args = p.parse_args() - metrics: list[str] | str = ( - args.metrics if args.metrics in ("all",) - else [m.strip() for m in args.metrics.split(",") if m.strip()] - ) + metrics: list[str] | str = (args.metrics if args.metrics in ("all", ) else + [m.strip() for m in args.metrics.split(",") if m.strip()]) evaluator = create_evaluator(metrics=metrics, device=args.device) sample: dict = {"video": load_video(args.video)} @@ -59,7 +59,10 @@ def main() -> None: evaluator.shutdown() print(json.dumps( - {name: r.score for name, r in results.items()}, + { + name: r.score + for name, r in results.items() + }, indent=2, )) diff --git a/examples/inference/gradio/local/gradio_local_demo.py b/examples/inference/gradio/local/gradio_local_demo.py index d3683e23d..3a76aabd5 100644 --- a/examples/inference/gradio/local/gradio_local_demo.py +++ b/examples/inference/gradio/local/gradio_local_demo.py @@ -8,15 +8,15 @@ from fastvideo.entrypoints.video_generator import VideoGenerator from fastvideo.api.sampling_param import SamplingParam from copy import deepcopy - MODEL_PATH_MAPPING = { "FastWan2.1-T2V-1.3B": "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", # "FastWan2.2-TI2V-5B-FullAttn": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", } + def create_timing_display(inference_time, total_time, stage_execution_times, num_frames): dit_denoising_time = f"{stage_execution_times[5]:.2f}s" if len(stage_execution_times) > 5 else "N/A" - + timing_html = f"""

⏱️ Timing Breakdown

@@ -47,7 +47,7 @@ def create_timing_display(inference_time, total_time, stage_execution_times, num
{total_time:.2f}s
""" - + if inference_time > 0: fps = num_frames / inference_time timing_html += f""" @@ -55,8 +55,10 @@ def create_timing_display(inference_time, total_time, stage_execution_times, num Generation Speed: {fps:.1f} frames/second """ - + return timing_html + "" + + def setup_model_environment(model_path: str) -> None: if "fullattn" in model_path.lower(): os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" @@ -64,10 +66,12 @@ def setup_model_environment(model_path: str) -> None: os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN" os.environ["FASTVIDEO_STAGE_LOGGING"] = "1" + def load_example_prompts(): + def contains_chinese(text): return any('\u4e00' <= char <= '\u9fff' for char in text) - + def load_from_file(filepath): prompts, labels = [], [] try: @@ -81,27 +85,28 @@ def load_example_prompts(): except Exception as e: print(f"Warning: Could not read {filepath}: {e}") return prompts, labels - + examples, example_labels = load_from_file("examples/inference/gradio/local/prompts_final.txt") - + if not examples: - examples = ["A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background."] + examples = [ + "A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background." + ] example_labels = ["Crowded rooftop bar at night"] - + return examples, example_labels def create_gradio_interface(default_params: dict[str, SamplingParam], generators: dict[str, VideoGenerator]): - def generate_video( - prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, - num_frames, height, width, randomize_seed, model_selection, progress - ): + + def generate_video(prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, num_frames, height, width, + randomize_seed, model_selection, progress): model_path = MODEL_PATH_MAPPING.get(model_selection, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers") setup_model_environment(model_path) try: if progress: progress(0.1, desc="Loading model for local inference...") - + generator = generators[model_path] params = deepcopy(default_params[model_path]) total_start_time = time.time() @@ -129,20 +134,25 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators output_dir = "outputs/" os.makedirs(output_dir, exist_ok=True) start_time = time.time() - result = generator.generate_video(prompt=prompt, sampling_param=params, save_video=True, return_frames=False) + result = generator.generate_video(prompt=prompt, + sampling_param=params, + save_video=True, + return_frames=False) inference_time = time.time() - start_time logging_info = result.get("logging_info", None) if logging_info: stage_names = logging_info.get_execution_order() stage_execution_times = [ - logging_info.get_stage_info(stage_name).get("execution_time", 0.0) - for stage_name in stage_names + logging_info.get_stage_info(stage_name).get("execution_time", 0.0) for stage_name in stage_names ] else: stage_names = [] stage_execution_times = [] total_time = time.time() - total_start_time - timing_details=create_timing_display(inference_time=inference_time, total_time=total_time, stage_execution_times=stage_execution_times, num_frames=params.num_frames) + timing_details = create_timing_display(inference_time=inference_time, + total_time=total_time, + stage_execution_times=stage_execution_times, + num_frames=params.num_frames) safe_prompt = params.prompt[:100].replace(' ', '_').replace('/', '_').replace('\\', '_') video_filename = f"{params.prompt[:100]}.mp4" output_path = os.path.join(output_dir, video_filename) @@ -157,7 +167,7 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators return None, f"Generation failed: {str(e)}", "" examples, example_labels = load_example_prompts() - + theme = gr.themes.Base().set( button_primary_background_fill="#2563eb", button_primary_background_fill_hover="#1d4ed8", @@ -165,7 +175,7 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators slider_color="#2563eb", checkbox_background_color_selected="#2563eb", ) - + def get_default_values(model_name): model_path = MODEL_PATH_MAPPING.get(model_name) if model_path and model_path in default_params: @@ -177,7 +187,7 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators 'guidance_scale': params.guidance_scale, 'seed': params.seed, } - + return { 'height': 448, 'width': 832, @@ -185,9 +195,9 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators 'guidance_scale': 3.0, 'seed': 1024, } - + initial_values = get_default_values("FastWan2.1-T2V-1.3B") - + with gr.Blocks(title="FastWan", theme=theme) as demo: gr.Image("assets/full.svg", show_label=False, container=False, height=80) @@ -197,7 +207,7 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators

Code | Blog | Docs

""") - + with gr.Accordion("🎥 What Is FastVideo?", open=False): gr.HTML("""
@@ -206,24 +216,20 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators

""") - - with gr.Row(): - model_selection = gr.Dropdown( - choices=list(MODEL_PATH_MAPPING.keys()), - value="FastWan2.1-T2V-1.3B", - label="Select Model", - interactive=True - ) with gr.Row(): - example_dropdown = gr.Dropdown( - choices=example_labels, - label="Example Prompts", - value=None, - interactive=True, - allow_custom_value=False - ) - + model_selection = gr.Dropdown(choices=list(MODEL_PATH_MAPPING.keys()), + value="FastWan2.1-T2V-1.3B", + label="Select Model", + interactive=True) + + with gr.Row(): + example_dropdown = gr.Dropdown(choices=example_labels, + label="Example Prompts", + value=None, + interactive=True, + allow_custom_value=False) + with gr.Row(): with gr.Column(scale=6): prompt = gr.Text( @@ -237,7 +243,7 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators ) with gr.Column(scale=1, min_width=120, elem_classes="center-button"): run_button = gr.Button("Run", variant="primary", size="lg") - + with gr.Row(): with gr.Column(): error_output = gr.Text(label="Error", visible=False) @@ -246,38 +252,32 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators with gr.Row(equal_height=True, elem_classes="main-content-row"): with gr.Column(scale=1, elem_classes="advanced-options-column"): with gr.Group(): - gr.HTML("
Advanced Options
") + gr.HTML( + "
Advanced Options
") with gr.Row(): - height = gr.Number( - label="Height", - value=initial_values['height'], - interactive=False, - container=True - ) - width = gr.Number( - label="Width", - value=initial_values['width'], - interactive=False, - container=True - ) - + height = gr.Number(label="Height", + value=initial_values['height'], + interactive=False, + container=True) + width = gr.Number(label="Width", + value=initial_values['width'], + interactive=False, + container=True) + with gr.Row(): - num_frames = gr.Number( - label="Number of Frames", - value=initial_values['num_frames'], - interactive=False, - container=True - ) + num_frames = gr.Number(label="Number of Frames", + value=initial_values['num_frames'], + interactive=False, + container=True) guidance_scale = gr.Slider( label="Guidance Scale", minimum=1, maximum=12, value=initial_values['guidance_scale'], ) - + with gr.Row(): - use_negative_prompt = gr.Checkbox( - label="Use negative prompt", value=False) + use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False) negative_prompt = gr.Text( label="Negative prompt", max_lines=3, @@ -295,17 +295,15 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators ) randomize_seed = gr.Checkbox(label="Randomize seed", value=False) seed_output = gr.Number(label="Used Seed") - + with gr.Column(scale=1, elem_classes="video-column"): - result = gr.Video( - label="Generated Video", - show_label=True, - height=466, - width=600, - container=True, - elem_classes="video-component" - ) - + result = gr.Video(label="Generated Video", + show_label=True, + height=466, + width=600, + container=True, + elem_classes="video-component") + gr.HTML(""" """) - + def on_example_select(example_label): if example_label and example_label in example_labels: index = example_labels.index(example_label) return examples[index] return "" - + example_dropdown.change( fn=on_example_select, inputs=example_dropdown, outputs=prompt, ) - + gr.HTML("""

Note that this demo is meant to showcase FastWan's quality and that under a large number of requests, generation speed may be affected.

""") - + use_negative_prompt.change( fn=lambda x: gr.update(visible=x), inputs=use_negative_prompt, outputs=negative_prompt, ) - + def on_model_selection_change(selected_model): if not selected_model: selected_model = "FastWan2.1-T2V-1.3B" - + model_path = MODEL_PATH_MAPPING.get(selected_model) - + if model_path and model_path in default_params: params = default_params[model_path] return ( @@ -476,7 +474,7 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators gr.update(value=params.guidance_scale), gr.update(value=params.seed), ) - + return ( gr.update(value=448), gr.update(value=832), @@ -484,35 +482,34 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators gr.update(value=3.0), gr.update(value=1024), ) - + model_selection.change( fn=on_model_selection_change, inputs=model_selection, outputs=[height, width, num_frames, guidance_scale, seed], ) - + def handle_generation(*args, progress=None, request: gr.Request = None): model_selection, prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, num_frames, height, width, randomize_seed = args - - result_path, seed_or_error, timing_details = generate_video( - prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, - num_frames, height, width, randomize_seed, model_selection, progress - ) + + result_path, seed_or_error, timing_details = generate_video(prompt, negative_prompt, use_negative_prompt, + seed, guidance_scale, num_frames, height, width, + randomize_seed, model_selection, progress) if result_path and os.path.exists(result_path): return ( - result_path, - seed_or_error, + result_path, + seed_or_error, gr.update(visible=False), gr.update(visible=True, value=timing_details), ) else: return ( - None, - seed_or_error, + None, + seed_or_error, gr.update(visible=True, value=seed_or_error), gr.update(visible=False), ) - + run_button.click( fn=handle_generation, inputs=[ @@ -530,19 +527,18 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators outputs=[result, seed_output, error_output, timing_display], concurrency_limit=20, ) - + return demo def main(): parser = argparse.ArgumentParser(description="FastVideo Gradio Local Demo") - parser.add_argument("--t2v_model_paths", type=str, + parser.add_argument("--t2v_model_paths", + type=str, default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers", help="Comma separated list of paths to the T2V model(s)") - parser.add_argument("--host", type=str, default="0.0.0.0", - help="Host to bind to") - parser.add_argument("--port", type=int, default=7860, - help="Port to bind to") + parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind to") + parser.add_argument("--port", type=int, default=7860, help="Port to bind to") args = parser.parse_args() generators = {} default_params = {} @@ -555,40 +551,36 @@ def main(): demo = create_gradio_interface(default_params, generators) print(f"Starting Gradio frontend at http://{args.host}:{args.port}") print(f"T2V Models: {args.t2v_model_paths}") - + from fastapi import FastAPI, Request, HTTPException from fastapi.responses import HTMLResponse, FileResponse import uvicorn - + app = FastAPI() - + @app.get("/logo.png") def get_logo(): - return FileResponse( - "assets/full.svg", - media_type="image/svg+xml", - headers={ - "Cache-Control": "public, max-age=3600", - "Access-Control-Allow-Origin": "*" - } - ) - + return FileResponse("assets/full.svg", + media_type="image/svg+xml", + headers={ + "Cache-Control": "public, max-age=3600", + "Access-Control-Allow-Origin": "*" + }) + @app.get("/favicon.ico") def get_favicon(): favicon_path = "assets/icon-simple.svg" - + if os.path.exists(favicon_path): - return FileResponse( - favicon_path, - media_type="image/svg+xml", - headers={ - "Cache-Control": "public, max-age=3600", - "Access-Control-Allow-Origin": "*" - } - ) + return FileResponse(favicon_path, + media_type="image/svg+xml", + headers={ + "Cache-Control": "public, max-age=3600", + "Access-Control-Allow-Origin": "*" + }) else: raise HTTPException(status_code=404, detail="Favicon not found") - + @app.get("/", response_class=HTMLResponse) def index(request: Request): base_url = str(request.base_url).rstrip('/') @@ -640,17 +632,16 @@ def main(): """ - - app = gr.mount_gradio_app( - app, - demo, - path="/gradio", - allowed_paths=[os.path.abspath("outputs"), os.path.abspath("fastvideo-logos")] - ) - + + app = gr.mount_gradio_app(app, + demo, + path="/gradio", + allowed_paths=[os.path.abspath("outputs"), + os.path.abspath("fastvideo-logos")]) + uvicorn.run(app, host=args.host, port=args.port) if __name__ == "__main__": - - main() \ No newline at end of file + + main() diff --git a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/app.py b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/app.py index 56ce123e4..32c1242bd 100644 --- a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/app.py +++ b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/app.py @@ -20,15 +20,15 @@ from .config import ( ) from .ui import create_gradio_interface + def main(): parser = argparse.ArgumentParser(description="FastVideo Gradio Local Demo") - parser.add_argument("--t2v_model_paths", type=str, + parser.add_argument("--t2v_model_paths", + type=str, default=MODEL_ID, help="Comma separated list of paths to the T2V model(s)") - parser.add_argument("--host", type=str, default="0.0.0.0", - help="Host to bind to") - parser.add_argument("--port", type=int, default=7860, - help="Port to bind to") + parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind to") + parser.add_argument("--port", type=int, default=7860, help="Port to bind to") args = parser.parse_args() gradio_temp_dir = os.path.abspath("outputs/gradio_tmp") os.makedirs(gradio_temp_dir, exist_ok=True) @@ -71,54 +71,46 @@ def main(): text_encoder_cpu_offload=False, ltx2_vae_tiling=False, ) - default_params[model_path] = apply_ltx2_defaults( - SamplingParam.from_pretrained(str(resolved_model_path)) - ) + default_params[model_path] = apply_ltx2_defaults(SamplingParam.from_pretrained(str(resolved_model_path))) demo = create_gradio_interface(default_params, generators) print(f"Starting Gradio frontend at http://{args.host}:{args.port}") print(f"T2V Models: {args.t2v_model_paths}") - + from fastapi import FastAPI, Request, HTTPException from fastapi.responses import HTMLResponse, FileResponse import uvicorn - + app = FastAPI() - + @app.get("/logo.png") def get_logo(): - return FileResponse( - "assets/full.svg", - media_type="image/svg+xml", - headers={ - "Cache-Control": "public, max-age=3600", - "Access-Control-Allow-Origin": "*" - } - ) + return FileResponse("assets/full.svg", + media_type="image/svg+xml", + headers={ + "Cache-Control": "public, max-age=3600", + "Access-Control-Allow-Origin": "*" + }) @app.get("/nvidia.png") def get_nvidia_logo(): - return FileResponse( - "assets/nv.png", - media_type="image/png", - headers={ - "Cache-Control": "public, max-age=3600", - "Access-Control-Allow-Origin": "*" - } - ) - + return FileResponse("assets/nv.png", + media_type="image/png", + headers={ + "Cache-Control": "public, max-age=3600", + "Access-Control-Allow-Origin": "*" + }) + @app.get("/favicon.ico") def get_favicon(): favicon_path = "assets/icon-simple.svg" - + if os.path.exists(favicon_path): - return FileResponse( - favicon_path, - media_type="image/svg+xml", - headers={ - "Cache-Control": "public, max-age=3600", - "Access-Control-Allow-Origin": "*" - } - ) + return FileResponse(favicon_path, + media_type="image/svg+xml", + headers={ + "Cache-Control": "public, max-age=3600", + "Access-Control-Allow-Origin": "*" + }) else: raise HTTPException(status_code=404, detail="Favicon not found") @@ -138,7 +130,7 @@ def main(): "Access-Control-Allow-Origin": "*", }, ) - + @app.get("/", response_class=HTMLResponse) def index(request: Request): base_url = str(request.base_url).rstrip('/') @@ -210,16 +202,14 @@ def main(): """ - - app = gr.mount_gradio_app( - app, - demo, - path="/gradio", - allowed_paths=[ - os.path.abspath("outputs"), - os.path.abspath("outputs_video"), - os.path.abspath("fastvideo-logos"), - ] - ) - + + app = gr.mount_gradio_app(app, + demo, + path="/gradio", + allowed_paths=[ + os.path.abspath("outputs"), + os.path.abspath("outputs_video"), + os.path.abspath("fastvideo-logos"), + ]) + uvicorn.run(app, host=args.host, port=args.port) diff --git a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/config.py b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/config.py index cff0e593b..febeb2135 100644 --- a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/config.py +++ b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/config.py @@ -7,19 +7,10 @@ import torch._inductor.config from fastvideo import SamplingParam LOCAL_DEMO_DIR = Path(__file__).resolve().parent -CLASSIFIER_DIR = Path( - os.path.expandvars( - os.path.expanduser( - os.getenv("LTX2_CLASSIFIER_DIR", str(LOCAL_DEMO_DIR)) - ) - ) -) +CLASSIFIER_DIR = Path(os.path.expandvars(os.path.expanduser(os.getenv("LTX2_CLASSIFIER_DIR", str(LOCAL_DEMO_DIR))))) MODEL_ID = os.path.expandvars( - os.path.expanduser( - os.getenv("LTX2_3_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers") - ) -) + os.path.expanduser(os.getenv("LTX2_3_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers"))) MODEL_PATH_MAPPING = { "FastLTX-2.3": MODEL_ID, } @@ -44,6 +35,7 @@ config.coordinate_descent_tuning = True config.coordinate_descent_check_all_directions = True config.epilogue_fusion = False + def apply_ltx2_defaults(params: SamplingParam) -> SamplingParam: params.height = DEFAULT_HEIGHT params.width = DEFAULT_WIDTH @@ -55,9 +47,11 @@ def apply_ltx2_defaults(params: SamplingParam) -> SamplingParam: params.negative_prompt = DEFAULT_NEGATIVE_PROMPT return params + def resolve_model_path(model_path: str) -> Path: return Path(os.path.expandvars(os.path.expanduser(model_path))) + def resolve_refine_upsampler_path(model_path: Path) -> Path: candidates = [ model_path / "spatial_upscaler", @@ -68,21 +62,18 @@ def resolve_refine_upsampler_path(model_path: Path) -> Path: env_path = os.getenv("LTX2_REFINE_UPSAMPLER_PATH") if env_path: - candidates.insert( - 0, Path(os.path.expandvars(os.path.expanduser(env_path))) - ) + candidates.insert(0, Path(os.path.expandvars(os.path.expanduser(env_path)))) for candidate in candidates: if (candidate / "config.json").is_file(): return candidate checked = "\n".join(f" - {candidate}" for candidate in candidates) - raise FileNotFoundError( - "Could not find an LTX2 refine upsampler directory.\n" - "Checked:\n" - f"{checked}\n" - "Set LTX2_REFINE_UPSAMPLER_PATH or update REFINE_UPSAMPLER_PATH." - ) + raise FileNotFoundError("Could not find an LTX2 refine upsampler directory.\n" + "Checked:\n" + f"{checked}\n" + "Set LTX2_REFINE_UPSAMPLER_PATH or update REFINE_UPSAMPLER_PATH.") + def setup_model_environment(model_path: str) -> None: _ = model_path diff --git a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/download_fasttext_classifiers.py b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/download_fasttext_classifiers.py index b1e484243..86a994c75 100644 --- a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/download_fasttext_classifiers.py +++ b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/download_fasttext_classifiers.py @@ -21,14 +21,16 @@ class ClassifierSpec: CLASSIFIERS = { - "nsfw": ClassifierSpec( + "nsfw": + ClassifierSpec( name="nsfw", repo_id="allenai/dolma-jigsaw-fasttext-bigrams-nsfw", source_filename="model.bin", target_filename="jigsaw_fasttext_bigrams_nsfw_final.bin", env_var="LTX2_NSFW_CLASSIFIER_PATH", ), - "hatespeech": ClassifierSpec( + "hatespeech": + ClassifierSpec( name="hatespeech", repo_id="allenai/dolma-jigsaw-fasttext-bigrams-hatespeech", source_filename="model.bin", @@ -48,20 +50,14 @@ def default_output_dir() -> Path: def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description=( - "Download the fastText safety classifiers used by " - "gradio_local_demo_ltx2_3.py." - ) - ) + parser = argparse.ArgumentParser(description=("Download the fastText safety classifiers used by " + "gradio_local_demo_ltx2_3.py.")) parser.add_argument( "--output-dir", type=Path, default=default_output_dir(), - help=( - "Directory where the demo looks for classifiers. Defaults to " - "LTX2_CLASSIFIER_DIR or the local Gradio demo directory." - ), + help=("Directory where the demo looks for classifiers. Defaults to " + "LTX2_CLASSIFIER_DIR or the local Gradio demo directory."), ) parser.add_argument( "--classifier", @@ -108,10 +104,8 @@ def download_classifier( print(f"[skip] {spec.name}: {destination}") return destination - print( - f"[download] {spec.name}: " - f"{spec.repo_id}/{spec.source_filename}" - ) + print(f"[download] {spec.name}: " + f"{spec.repo_id}/{spec.source_filename}") cached_path = Path( hf_hub_download( repo_id=spec.repo_id, @@ -119,8 +113,7 @@ def download_classifier( cache_dir=cache_dir, force_download=force, token=token, - ) - ) + )) materialize_download(cached_path, destination) print(f"[ready] {spec.name}: {destination}") return destination @@ -138,9 +131,7 @@ def print_summary( print("\nThese filenames are auto-discovered by gradio_local_demo_ltx2_3.py.") if output_dir != default_output_dir(): - print( - "\nBecause you used a custom output directory, point the demo to it:" - ) + print("\nBecause you used a custom output directory, point the demo to it:") print(f' export LTX2_CLASSIFIER_DIR="{output_dir}"') for spec, path in zip(selected_specs, downloaded_paths, strict=True): @@ -164,8 +155,7 @@ def main() -> None: cache_dir=cache_dir, force=args.force, token=args.token, - ) - for spec in selected_specs + ) for spec in selected_specs ] print_summary(selected_specs, output_dir, downloaded_paths) diff --git a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/examples.py b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/examples.py index 50147d4eb..57b2f1cef 100644 --- a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/examples.py +++ b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/examples.py @@ -1,6 +1,7 @@ import json import os + def load_example_prompts(): examples: list[str] = [] example_labels: list[str] = [] @@ -23,9 +24,7 @@ def load_example_prompts(): if not prompt: continue examples.append(prompt) - example_labels.append( - prompt[:100] + "..." if len(prompt) > 100 else prompt - ) + example_labels.append(prompt[:100] + "..." if len(prompt) > 100 else prompt) except Exception as e: print(f"Warning: Could not read {prompts_path}: {e}") @@ -52,9 +51,7 @@ def load_example_prompts(): if not caption: continue examples.append(caption) - example_labels.append( - caption[:100] + "..." if len(caption) > 100 else caption - ) + example_labels.append(caption[:100] + "..." if len(caption) > 100 else caption) except Exception as e: print(f"Warning: Could not read {validation_path}: {e}") diff --git a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/prompt_enhancer.py b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/prompt_enhancer.py index 7ba13adf8..e9e3dfb94 100644 --- a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/prompt_enhancer.py +++ b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/prompt_enhancer.py @@ -14,9 +14,7 @@ except ImportError: # pragma: no cover - optional dependency Cerebras = None # type: ignore[assignment] LOCAL_DEMO_DIR = Path(__file__).resolve().parent -DEFAULT_SYSTEM_PROMPT_PATH = ( - LOCAL_DEMO_DIR / "prompts" / "prompt_extension_system_prompt.md" -) +DEFAULT_SYSTEM_PROMPT_PATH = (LOCAL_DEMO_DIR / "prompts" / "prompt_extension_system_prompt.md") DEFAULT_PROVIDER = "cerebras" DEFAULT_MODEL = "gpt-oss-120b" DEFAULT_TEMPERATURE = 1.0 @@ -83,19 +81,13 @@ def _load_prompt_required(path: str, prompt_name: str) -> str: try: text = candidate.read_text(encoding="utf-8").strip() except OSError as exc: - raise RuntimeError( - f"Failed to read {prompt_name} system prompt: {candidate}" - ) from exc + raise RuntimeError(f"Failed to read {prompt_name} system prompt: {candidate}") from exc if not text: - raise RuntimeError( - f"{prompt_name} system prompt file is empty: {candidate}" - ) + raise RuntimeError(f"{prompt_name} system prompt file is empty: {candidate}") return text tried = ", ".join(str(candidate) for candidate in candidate_paths) - raise RuntimeError( - f"{prompt_name} system prompt file not found. Tried: {tried}" - ) + raise RuntimeError(f"{prompt_name} system prompt file not found. Tried: {tried}") def _extract_assistant_content(response_json: dict[str, Any]) -> str: @@ -129,10 +121,8 @@ def _extract_assistant_content(response_json: dict[str, Any]) -> str: finish_reason = choices[0].get("finish_reason") refusal = message.get("refusal") - raise ValueError( - "Missing assistant content in chat completion response. " - f"finish_reason={finish_reason!r}, refusal={refusal!r}" - ) + raise ValueError("Missing assistant content in chat completion response. " + f"finish_reason={finish_reason!r}, refusal={refusal!r}") def _dump_response_json(response: Any) -> dict[str, Any]: @@ -143,15 +133,11 @@ def _dump_response_json(response: Any) -> dict[str, Any]: elif hasattr(response, "dict"): payload = response.dict() else: - raise TypeError( - "Unsupported chat completion response type. " - f"type={type(response)!r}" - ) + raise TypeError("Unsupported chat completion response type. " + f"type={type(response)!r}") if not isinstance(payload, dict): - raise TypeError( - "Chat completion response did not serialize to a JSON object." - ) + raise TypeError("Chat completion response did not serialize to a JSON object.") return payload @@ -203,6 +189,7 @@ def _normalize_prompt(value: Any) -> str: class PromptEnhancer: + def __init__(self) -> None: self.provider = DEFAULT_PROVIDER self.provider_label = DEFAULT_PROVIDER @@ -215,9 +202,8 @@ class PromptEnhancer: "LTX2_PROMPT_TEMPERATURE", DEFAULT_TEMPERATURE, ) - self.system_prompt_path = _env_value( - "LTX2_PROMPT_EXTENSION_SYSTEM_PROMPT_PATH" - ) or str(DEFAULT_SYSTEM_PROMPT_PATH) + self.system_prompt_path = _env_value("LTX2_PROMPT_EXTENSION_SYSTEM_PROMPT_PATH") or str( + DEFAULT_SYSTEM_PROMPT_PATH) self.client: Any | None = None self.system_prompt: str | None = None self.unavailable_reason: str | None = None @@ -237,14 +223,10 @@ class PromptEnhancer: def _build_client(self) -> Any: if Cerebras is None: - raise RuntimeError( - "Cerebras SDK is not installed. Install " - "'cerebras-cloud-sdk' to use prompt enhancement." - ) + raise RuntimeError("Cerebras SDK is not installed. Install " + "'cerebras-cloud-sdk' to use prompt enhancement.") if not self.api_key: - raise RuntimeError( - "Missing FASTVIDEO_PROMPT_API_KEY or CEREBRAS_API_KEY." - ) + raise RuntimeError("Missing FASTVIDEO_PROMPT_API_KEY or CEREBRAS_API_KEY.") return Cerebras(api_key=self.api_key) def _build_body( @@ -260,11 +242,14 @@ class PromptEnhancer: "single 5-second LTX-2.3 video. Respond with valid JSON " 'only as {"prompt": "..."}.' # noqa: E501 ), - "user_prompt": user_prompt, + "user_prompt": + user_prompt, } return { - "model": model or self.model, - "temperature": self.temperature, + "model": + model or self.model, + "temperature": + self.temperature, "messages": [ { "role": "system", @@ -285,10 +270,7 @@ class PromptEnhancer: model: str | None = None, ) -> tuple[dict[str, Any], str]: if self.client is None: - raise RuntimeError( - self.unavailable_reason - or "Prompt enhancement client is not initialized." - ) + raise RuntimeError(self.unavailable_reason or "Prompt enhancement client is not initialized.") body = self._build_body( system_prompt=system_prompt, @@ -332,10 +314,7 @@ class PromptEnhancer: return EnhanceResult( prompt="", fallback_used=True, - error=( - self.unavailable_reason - or "Prompt enhancement system prompt is unavailable." - ), + error=(self.unavailable_reason or "Prompt enhancement system prompt is unavailable."), provider=self.provider_label, model=self.model, latency_ms=0.0, @@ -363,10 +342,8 @@ class PromptEnhancer: except Exception as exc: error_detail = str(exc) if isinstance(response_content, str) and response_content.strip(): - error_detail = ( - f"{error_detail} | assistant_response=" - f"{_preview_text(response_content, limit=240)}" - ) + error_detail = (f"{error_detail} | assistant_response=" + f"{_preview_text(response_content, limit=240)}") latency_ms = (time.perf_counter() - t0) * 1000.0 return EnhanceResult( prompt="", diff --git a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/prompt_rewrite.py b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/prompt_rewrite.py index fce831c39..4bc744fa5 100644 --- a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/prompt_rewrite.py +++ b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/prompt_rewrite.py @@ -2,10 +2,12 @@ from functools import lru_cache from .prompt_enhancer import PromptEnhancer + @lru_cache(maxsize=1) def get_prompt_enhancer() -> PromptEnhancer: return PromptEnhancer() + def maybe_enhance_prompt( prompt: str, curated_prompts: set[str], @@ -15,28 +17,20 @@ def maybe_enhance_prompt( return normalized_prompt if normalized_prompt in curated_prompts: - print( - "[ENHANCE][INFO] Skipping prompt enhancement for curated prompt." - ) + print("[ENHANCE][INFO] Skipping prompt enhancement for curated prompt.") return normalized_prompt enhancer = get_prompt_enhancer() if not enhancer.api_key: - raise RuntimeError( - "Prompt enhancement is enabled for custom prompts, but " - "FASTVIDEO_PROMPT_API_KEY or CEREBRAS_API_KEY is not set." - ) + raise RuntimeError("Prompt enhancement is enabled for custom prompts, but " + "FASTVIDEO_PROMPT_API_KEY or CEREBRAS_API_KEY is not set.") result = enhancer.enhance_prompt(normalized_prompt) if result.fallback_used or not result.prompt.strip(): - print( - "[ENHANCE][WARN] Falling back to raw prompt " - f"error={result.error}" - ) + print("[ENHANCE][WARN] Falling back to raw prompt " + f"error={result.error}") return normalized_prompt - print( - "[ENHANCE][INFO] Prompt enhanced " - f"latency={result.latency_ms:.2f}ms model={result.model}" - ) + print("[ENHANCE][INFO] Prompt enhanced " + f"latency={result.latency_ms:.2f}ms model={result.model}") return result.prompt.strip() diff --git a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/rendering.py b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/rendering.py index 18b931f4c..e4ef0d1a3 100644 --- a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/rendering.py +++ b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/rendering.py @@ -3,6 +3,7 @@ from pathlib import Path from .config import DEFAULT_FPS, GENERATED_CLIP_ROOT, MAX_SESSION_CLIPS + def create_timing_display(inference_time, total_time, stage_execution_times, num_frames): timing_html = f"""
@@ -18,7 +19,7 @@ def create_timing_display(inference_time, total_time, stage_execution_times, num
{total_time:.1f}s
""" - + if inference_time > 0: fps = num_frames / inference_time timing_html += f""" @@ -26,9 +27,10 @@ def create_timing_display(inference_time, total_time, stage_execution_times, num Generation Speed: {fps:.1f} frames/second """ - + return timing_html + "" + def create_timing_placeholder() -> str: return """
@@ -49,19 +51,23 @@ def create_timing_placeholder() -> str:
""" + def _truncate_text(value: str, max_chars: int) -> str: if len(value) <= max_chars: return value - return value[: max_chars - 1].rstrip() + "..." + return value[:max_chars - 1].rstrip() + "..." + def _clip_duration_seconds(num_frames: int, fps: int) -> int: return max(1, round(num_frames / max(fps, 1))) + def _make_clip_public_path(output_path: str) -> str: resolved_path = Path(output_path).resolve() relative_path = resolved_path.relative_to(GENERATED_CLIP_ROOT.resolve()) return f"/generated-clips/{relative_path.as_posix()}" + def _record_session_clip( session_clips: list[dict[str, str | int | float]], *, @@ -86,6 +92,7 @@ def _record_session_clip( updated_clips = updated_clips[-MAX_SESSION_CLIPS:] return updated_clips + def render_completed_clips(clips: list[dict[str, str | int | float]]) -> str: if not clips: return """ @@ -104,8 +111,7 @@ def render_completed_clips(clips: list[dict[str, str | int | float]]) -> str: prompt_preview = html.escape(str(clip["prompt_preview"])) model_name = html.escape(str(clip["model_name"])) duration_label = html.escape(str(clip["duration_label"])) - cards.append( - f""" + cards.append(f"""
@@ -122,8 +128,7 @@ def render_completed_clips(clips: list[dict[str, str | int | float]]) -> str:
- """ - ) + """) return f"""
@@ -131,6 +136,7 @@ def render_completed_clips(clips: list[dict[str, str | int | float]]) -> str:
""" + def render_error_message(message: str) -> str: return f"""
@@ -139,17 +145,16 @@ def render_error_message(message: str) -> str:
""" + def render_prompt_blocked_message( message: str, category: str | None = None, ) -> str: details = "" if category: - details = ( - '
' - f"Policy: {html.escape(category)}" - "
" - ) + details = ('
' + f"Policy: {html.escape(category)}" + "
") return f"""
Prompt Blocked
@@ -158,13 +163,12 @@ def render_prompt_blocked_message(
""" + def render_input_image_status(input_image: str | None) -> str: if not input_image: return "" image_name = html.escape(Path(str(input_image)).name) - return ( - "
" - f"Image ready: {image_name}" - "
" - ) + return ("
" + f"Image ready: {image_name}" + "
") diff --git a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/safety.py b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/safety.py index 66010eb7b..2c41de7da 100644 --- a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/safety.py +++ b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/safety.py @@ -8,6 +8,7 @@ import fasttext from .config import CLASSIFIER_DIR + def resolve_classifier_path( classifier_kind: str, env_var: str, @@ -18,29 +19,24 @@ def resolve_classifier_path( candidates: list[Path] = [] env_path = os.getenv(env_var) if env_path: - candidates.append( - Path(os.path.expandvars(os.path.expanduser(env_path))) - ) - candidates.extend( - [ - CLASSIFIER_DIR / filename, - Path(f"/home/shared/{shared_filename}"), - Path(legacy_path), - ] - ) + candidates.append(Path(os.path.expandvars(os.path.expanduser(env_path)))) + candidates.extend([ + CLASSIFIER_DIR / filename, + Path(f"/home/shared/{shared_filename}"), + Path(legacy_path), + ]) for candidate in candidates: if candidate.is_file(): return str(candidate) checked = "\n".join(f" - {candidate}" for candidate in candidates) - raise FileNotFoundError( - f"Could not find the {classifier_kind} classifier.\n" - f"Checked:\n{checked}\n" - "Run " - "`python examples/inference/gradio/local/download_fasttext_classifiers.py` " - "or set the appropriate classifier path environment variable." - ) + raise FileNotFoundError(f"Could not find the {classifier_kind} classifier.\n" + f"Checked:\n{checked}\n" + "Run " + "`python examples/inference/gradio/local/download_fasttext_classifiers.py` " + "or set the appropriate classifier path environment variable.") + def fasttext_predict( model_path: str, @@ -62,17 +58,17 @@ def fasttext_predict( latency_ms = (time.perf_counter() - start_time) * 1000.0 identifier = labels[0].replace('__label__', '') confidence = probs[0] - print( - "[safety] " - f"{classifier_name} fastText latency={latency_ms:.2f}ms " - f"label={identifier} confidence={float(confidence):.4f}" - ) + print("[safety] " + f"{classifier_name} fastText latency={latency_ms:.2f}ms " + f"label={identifier} confidence={float(confidence):.4f}") return identifier, confidence + @lru_cache(maxsize=None) def load_fasttext_model(model_path: str): return fasttext.load_model(model_path) + def classify_nsfw(text: str) -> tuple[str, float]: return fasttext_predict( resolve_classifier_path( @@ -86,6 +82,7 @@ def classify_nsfw(text: str) -> tuple[str, float]: "nsfw", ) + def classify_toxic_speech(text: str) -> tuple[str, float]: return fasttext_predict( resolve_classifier_path( @@ -99,9 +96,11 @@ def classify_toxic_speech(text: str) -> tuple[str, float]: "hate_speech", ) + def _normalize_classifier_label(identifier: str) -> str: return identifier.strip().lower().replace("-", "_").replace(" ", "_") + def _label_matches( identifier: str, blocked_markers: tuple[str, ...], @@ -112,12 +111,14 @@ def _label_matches( return False return any(marker in normalized for marker in blocked_markers) + @dataclass(frozen=True) class PromptSafetyCheck: blocked: bool category: str | None = None message: str | None = None + def get_prompt_safety_check(prompt: str) -> PromptSafetyCheck: normalized_prompt = prompt.strip() if not normalized_prompt: @@ -125,41 +126,37 @@ def get_prompt_safety_check(prompt: str) -> PromptSafetyCheck: nsfw_label, _ = classify_nsfw(normalized_prompt) if _label_matches( - nsfw_label, - blocked_markers=("nsfw",), - safe_markers=("sfw", "safe"), + nsfw_label, + blocked_markers=("nsfw", ), + safe_markers=("sfw", "safe"), ): return PromptSafetyCheck( blocked=True, category="NSFW", - message=( - "This request was blocked by the safety filter because it " - "appears to contain NSFW content. Please revise the prompt " - "and try again." - ), + message=("This request was blocked by the safety filter because it " + "appears to contain NSFW content. Please revise the prompt " + "and try again."), ) hate_label, _ = classify_toxic_speech(normalized_prompt) if _label_matches( - hate_label, - blocked_markers=("hatespeech", "hate", "toxic", "offensive", "abusive"), - safe_markers=( - "non_hatespeech", - "not_hatespeech", - "non_toxic", - "not_toxic", - "safe", - "clean", - ), + hate_label, + blocked_markers=("hatespeech", "hate", "toxic", "offensive", "abusive"), + safe_markers=( + "non_hatespeech", + "not_hatespeech", + "non_toxic", + "not_toxic", + "safe", + "clean", + ), ): return PromptSafetyCheck( blocked=True, category="Hate Speech", - message=( - "This request was blocked by the safety filter because it " - "appears to contain hate speech or abusive content. Please " - "revise the prompt and try again." - ), + message=("This request was blocked by the safety filter because it " + "appears to contain hate speech or abusive content. Please " + "revise the prompt and try again."), ) return PromptSafetyCheck(blocked=False) diff --git a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/ui.py b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/ui.py index 1d08b1f95..6b25744ae 100644 --- a/examples/inference/gradio/local/gradio_local_demo_ltx2_3/ui.py +++ b/examples/inference/gradio/local/gradio_local_demo_ltx2_3/ui.py @@ -33,16 +33,16 @@ from .rendering import ( ) from .safety import get_prompt_safety_check + def create_gradio_interface(default_params: dict[str, SamplingParam], generators: dict[str, VideoGenerator]): + def _sanitize_filename_component(name: str) -> str: sanitized = re.sub(r'[\\/:*?"<>|]', "", name) sanitized = sanitized.strip().strip(".") sanitized = re.sub(r"\s+", "_", sanitized) return sanitized or "video" - def generate_video( - prompt, model_selection, input_image=None - ): + def generate_video(prompt, model_selection, input_image=None): model_path = MODEL_PATH_MAPPING.get(model_selection, MODEL_ID) setup_model_environment(model_path) try: @@ -69,45 +69,33 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators output_path = str(OUTPUT_DIR / video_filename) params.output_path = output_path start_time = time.perf_counter() - result = generator.generate_video( - prompt=prompt, - output_path=output_path, - fps=DEFAULT_FPS, - seed=int(params.seed), - save_video=True, - return_frames=False, - guidance_scale=float(params.guidance_scale), - height=int(params.height), - width=int(params.width), - num_frames=int(params.num_frames), - num_inference_steps=DEFAULT_NUM_INFERENCE_STEPS, - negative_prompt=params.negative_prompt, - image_path=params.image_path, - ltx2_image_crf=0.0 - ) + result = generator.generate_video(prompt=prompt, + output_path=output_path, + fps=DEFAULT_FPS, + seed=int(params.seed), + save_video=True, + return_frames=False, + guidance_scale=float(params.guidance_scale), + height=int(params.height), + width=int(params.width), + num_frames=int(params.num_frames), + num_inference_steps=DEFAULT_NUM_INFERENCE_STEPS, + negative_prompt=params.negative_prompt, + image_path=params.image_path, + ltx2_image_crf=0.0) wall_time = time.perf_counter() - start_time - generation_time = ( - result.get("generation_time") - if isinstance(result, dict) else None - ) - e2e_latency = ( - result.get("e2e_latency") - if isinstance(result, dict) else None - ) + generation_time = (result.get("generation_time") if isinstance(result, dict) else None) + e2e_latency = (result.get("e2e_latency") if isinstance(result, dict) else None) if generation_time is None: generation_time = wall_time if e2e_latency is None: e2e_latency = wall_time - resolved_output_path = ( - result.get("output_path", output_path) - if isinstance(result, dict) else output_path - ) + resolved_output_path = (result.get("output_path", output_path) if isinstance(result, dict) else output_path) logging_info = result.get("logging_info", None) if isinstance(result, dict) else None if logging_info: stage_names = logging_info.get_execution_order() stage_execution_times = [ - logging_info.get_stage_info(stage_name).get("execution_time", 0.0) - for stage_name in stage_names + logging_info.get_stage_info(stage_name).get("execution_time", 0.0) for stage_name in stage_names ] else: stage_names = [] @@ -126,11 +114,9 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators return None, f"Generation failed: {str(e)}", 0, 0.0, 0.0 examples, example_labels = load_example_prompts() - curated_prompts = { - prompt.strip() for prompt in examples if prompt.strip() - } + curated_prompts = {prompt.strip() for prompt in examples if prompt.strip()} initial_example_label = None - + theme = gr.themes.Base().set( button_primary_background_fill="#2563eb", button_primary_background_fill_hover="#1d4ed8", @@ -169,7 +155,8 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators
Real-Time 1080p Video Generation with FastLTX-2.3 on a single B200
- """, elem_id="hero-wrapper") + """, + elem_id="hero-wrapper") with gr.Column(elem_id="app-shell", elem_classes="app-shell"): timing_title = gr.HTML( @@ -186,9 +173,7 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators with gr.Group(elem_id="stage-card", elem_classes="stage-card"): with gr.Row(elem_id="stage-card-header", elem_classes="stage-card-header"): - gr.HTML( - "
🏎️ Make Video Generation Go Blurrrrrrr 💨
" - ) + gr.HTML("
🏎️ Make Video Generation Go Blurrrrrrr 💨
") stage_badges = gr.HTML( render_generation_badges("FastLTX-2.3"), elem_id="stage-badges", @@ -246,8 +231,8 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators elem_id="input-image", ) with gr.Row( - elem_id="image-upload-status-row", - elem_classes="image-upload-status-row", + elem_id="image-upload-status-row", + elem_classes="image-upload-status-row", ): image_upload_status = gr.HTML( value=render_input_image_status(None), @@ -334,7 +319,10 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators with gr.Row(elem_id="completed-clips-header-row", elem_classes="completed-clips-header-row"): with gr.Column(scale=4): gr.Markdown("## Gallery") - with gr.Column(scale=1, min_width=170, elem_id="completed-clips-button-column", elem_classes="completed-clips-button-column"): + with gr.Column(scale=1, + min_width=170, + elem_id="completed-clips-button-column", + elem_classes="completed-clips-button-column"): clear_clips_button = gr.Button( "Clear My Gallery", variant="secondary", @@ -1528,13 +1516,13 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators })(); """) - + def on_example_select(example_label): if example_label and example_label in example_labels: index = example_labels.index(example_label) return examples[index] return gr.update() - + example_dropdown.change( fn=on_example_select, inputs=example_dropdown, @@ -1586,7 +1574,7 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators stage_badges, ], ) - + def summarize_clip_status(session_clips): session_clips = session_clips or [] count = len(session_clips) @@ -1678,15 +1666,11 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators ) if prompt_for_generation != normalized_prompt: - enhanced_safety_check = get_prompt_safety_check( - prompt_for_generation - ) + enhanced_safety_check = get_prompt_safety_check(prompt_for_generation) if enhanced_safety_check.blocked: - message = ( - "Prompt enhancement produced text that was blocked by " - "the safety filter. Please revise the prompt and try " - "again." - ) + message = ("Prompt enhancement produced text that was blocked by " + "the safety filter. Please revise the prompt and try " + "again.") gr.Warning(message) return ( gr.update(value=None, visible=True), @@ -1708,8 +1692,7 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators ) result_path, seed_or_error, num_frames, generation_time, e2e_latency = generate_video( - prompt_for_generation, model_selection, input_image - ) + prompt_for_generation, model_selection, input_image) timing_details = create_timing_display( inference_time=generation_time, total_time=e2e_latency, @@ -1762,7 +1745,7 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators ], queue=False, ) - + run_button.click( fn=handle_generation, inputs=[ @@ -1784,5 +1767,5 @@ def create_gradio_interface(default_params: dict[str, SamplingParam], generators show_progress_on=result, queue=False, ) - + return demo diff --git a/examples/inference/gradio/local/gradio_local_demo_matrixgame2.py b/examples/inference/gradio/local/gradio_local_demo_matrixgame2.py index dc88f434f..f39c1b81b 100644 --- a/examples/inference/gradio/local/gradio_local_demo_matrixgame2.py +++ b/examples/inference/gradio/local/gradio_local_demo_matrixgame2.py @@ -12,32 +12,40 @@ from fastapi.responses import HTMLResponse, FileResponse from fastvideo.entrypoints.streaming_generator import StreamingVideoGenerator from fastvideo.models.dits.matrixgame2.utils import expand_action_to_frames - VARIANT_CONFIG = { "Matrix-Game-2.0-Base": { - "model_path": "FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers", - "keyboard_dim": 4, - "mode": "universal", - "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png", + "model_path": + "FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers", + "keyboard_dim": + 4, + "mode": + "universal", + "image_url": + "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png", }, "Matrix-Game-2.0-GTA": { - "model_path": "FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers", - "keyboard_dim": 2, - "mode": "gta_drive", - "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/gta_drive/0000.png", + "model_path": + "FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers", + "keyboard_dim": + 2, + "mode": + "gta_drive", + "image_url": + "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/gta_drive/0000.png", }, "Matrix-Game-2.0-TempleRun": { - "model_path": "FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers", - "keyboard_dim": 7, - "mode": "templerun", - "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/temple_run/0000.png", + "model_path": + "FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers", + "keyboard_dim": + 7, + "mode": + "templerun", + "image_url": + "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/temple_run/0000.png", }, } -MODEL_PATH_MAPPING = { - name: config["model_path"] for name, config in VARIANT_CONFIG.items() -} - +MODEL_PATH_MAPPING = {name: config["model_path"] for name, config in VARIANT_CONFIG.items()} CAM_VALUE = 0.1 KEYBOARD_MAP_UNIVERSAL = { @@ -62,7 +70,6 @@ KEYBOARD_MAP_TEMPLERUN = { "D (Right)": [0, 0, 0, 0, 0, 0, 1], } - CAMERA_MAP_UNIVERSAL = { "U (Center)": [0, 0], "I (Up)": [CAM_VALUE, 0], @@ -76,6 +83,7 @@ CAMERA_MAP_GTA = { "D (Steer Right)": [0, CAM_VALUE], } + def setup_model_environment(model_path: str) -> None: # if "fullattn" in model_path.lower(): # os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" @@ -84,9 +92,10 @@ def setup_model_environment(model_path: str) -> None: os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" os.environ["FASTVIDEO_STAGE_LOGGING"] = "1" + def create_timing_display(inference_time, total_time, stage_execution_times, num_frames): dit_denoising_time = f"{stage_execution_times[5]:.2f}s" if len(stage_execution_times) > 5 else "N/A" - + timing_html = f"""

⏱️ Timing Breakdown

@@ -117,7 +126,7 @@ def create_timing_display(inference_time, total_time, stage_execution_times, num
{total_time:.2f}s
""" - + if inference_time > 0: fps = num_frames / inference_time timing_html += f""" @@ -125,9 +134,10 @@ def create_timing_display(inference_time, total_time, stage_execution_times, num Generation Speed: {fps:.1f} frames/second """ - + return timing_html + "" + def get_action_tensors(mode: str, keyboard_key: str, mouse_key: str | None): if mode == "universal": keyboard = torch.tensor(KEYBOARD_MAP_UNIVERSAL.get(keyboard_key, [0, 0, 0, 0])).cuda() @@ -140,13 +150,14 @@ def get_action_tensors(mode: str, keyboard_key: str, mouse_key: str | None): mouse = None else: raise ValueError(f"Unknown mode: {mode}") - + return {"keyboard": keyboard, "mouse": mouse} + def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], loaded_model_name: str): initial_config = VARIANT_CONFIG.get(loaded_model_name, VARIANT_CONFIG["Matrix-Game-2.0-Base"]) initial_mode = initial_config["mode"] - + if initial_mode == "universal": initial_kb_choices = list(KEYBOARD_MAP_UNIVERSAL.keys()) initial_mouse_choices = list(CAMERA_MAP_UNIVERSAL.keys()) @@ -159,7 +170,7 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load initial_kb_choices = list(KEYBOARD_MAP_TEMPLERUN.keys()) initial_mouse_choices = [] initial_mouse_visible = False - + theme = gr.themes.Base().set( button_primary_background_fill="#2563eb", button_primary_background_fill_hover="#1d4ed8", @@ -167,7 +178,7 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load slider_color="#2563eb", checkbox_background_color_selected="#2563eb", ) - + with gr.Blocks(title="FastVideo - Matrix Game 2.0", theme=theme) as demo: game_state = gr.State({ "initialized": False, @@ -175,17 +186,17 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load "block_idx": 0, "max_blocks": 50, }) - + # Header gr.Image("assets/full.svg", show_label=False, container=False, height=80) - + gr.HTML("""

Make Video Generation Go Blurrrrrrr

Code | Blog | Docs

""") - + with gr.Accordion("🎥 What Is FastVideo?", open=False): gr.HTML("""
@@ -194,49 +205,46 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load

""") - + # Model Selection with gr.Row(): - model_selection = gr.Dropdown( - choices=[loaded_model_name], - value=loaded_model_name, - label="Select Model", - interactive=False - ) - + model_selection = gr.Dropdown(choices=[loaded_model_name], + value=loaded_model_name, + label="Select Model", + interactive=False) # Main Layout with gr.Row(equal_height=True, elem_classes="main-content-row"): with gr.Column(scale=1, elem_classes="advanced-options-column"): with gr.Group(): gr.HTML("
Game Controls
") - + with gr.Group(): - gr.HTML("
🎮 Keyboard Control
") - keyboard_action = gr.Radio( - choices=initial_kb_choices, - value=initial_kb_choices[0] if initial_kb_choices else None, - label="Movement", - show_label=False, - interactive=True + gr.HTML( + "
🎮 Keyboard Control
" ) - + keyboard_action = gr.Radio(choices=initial_kb_choices, + value=initial_kb_choices[0] if initial_kb_choices else None, + label="Movement", + show_label=False, + interactive=True) + with gr.Group(visible=initial_mouse_visible) as mouse_group: - gr.HTML("
🖱️ Mouse/Camera Control
") - mouse_action = gr.Radio( - choices=initial_mouse_choices if initial_mouse_visible else [], - value=initial_mouse_choices[0] if initial_mouse_choices else None, - label="Camera", - show_label=False, - interactive=True + gr.HTML( + "
🖱️ Mouse/Camera Control
" ) - + mouse_action = gr.Radio(choices=initial_mouse_choices if initial_mouse_visible else [], + value=initial_mouse_choices[0] if initial_mouse_choices else None, + label="Camera", + show_label=False, + interactive=True) + with gr.Row(): action_btn = gr.Button("Start", variant="primary") stop_btn = gr.Button("Stop", variant="stop") - + gr.HTML("
") - + seed = gr.Slider( label="Seed", minimum=0, @@ -246,21 +254,18 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load ) randomize_seed = gr.Checkbox(label="Randomize seed", value=False) seed_output = gr.Number(label="Used Seed") - - block_counter = gr.Textbox(label="Progress", value="Block: 0 / 50", interactive=False, lines=1) + block_counter = gr.Textbox(label="Progress", value="Block: 0 / 50", interactive=False, lines=1) # Right Column: Video Output with gr.Column(scale=1, elem_classes="video-column"): - video_output = gr.Video( - label="Generated Video", - show_label=True, - height=466, - width=600, - container=True, - elem_classes="video-component", - autoplay=True - ) + video_output = gr.Video(label="Generated Video", + show_label=True, + height=466, + width=600, + container=True, + elem_classes="video-component", + autoplay=True) # Styles gr.HTML(""" @@ -393,12 +398,12 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load } """) - + # UI update based on model selection def on_model_change(model_name): config = VARIANT_CONFIG.get(model_name, VARIANT_CONFIG["Matrix-Game-2.0-Base"]) mode = config["mode"] - + if mode == "universal": kb_choices = list(KEYBOARD_MAP_UNIVERSAL.keys()) mouse_choices = list(CAMERA_MAP_UNIVERSAL.keys()) @@ -411,31 +416,31 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load kb_choices = list(KEYBOARD_MAP_TEMPLERUN.keys()) mouse_choices = [] mouse_visible = False - + return ( gr.update(choices=kb_choices, value=kb_choices[0] if kb_choices else None), - gr.update(choices=mouse_choices, value=mouse_choices[0] if mouse_choices else None, visible=mouse_visible), + gr.update(choices=mouse_choices, + value=mouse_choices[0] if mouse_choices else None, + visible=mouse_visible), gr.update(visible=mouse_visible), ) - - model_selection.change( - fn=on_model_change, - inputs=model_selection, - outputs=[keyboard_action, mouse_action, mouse_group] - ) + + model_selection.change(fn=on_model_change, + inputs=model_selection, + outputs=[keyboard_action, mouse_action, mouse_group]) def start_game(model_name, seed_val, randomize, state): if randomize: - seed_val = torch.randint(0, 1000000, (1,)).item() - + seed_val = torch.randint(0, 1000000, (1, )).item() + config = VARIANT_CONFIG.get(model_name) if not config: return state, seed_val, "Block: 0 / 50", None, "", gr.update(), gr.update() - + generator = generators.get(config["model_path"]) if not generator: return state, seed_val, "Block: 0 / 50", None, "", gr.update(), gr.update() - + # If already initialized, clean up first if state.get("initialized"): try: @@ -444,23 +449,23 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load generator.executor.execute_streaming_clear() except Exception as e: print(f"Warning: cleanup error: {e}") - + # Streaming parameters num_latent_frames_per_block = 3 max_blocks = 50 total_latent_frames = num_latent_frames_per_block * max_blocks num_frames = (total_latent_frames - 1) * 4 + 1 - + actions = { "keyboard": torch.zeros((num_frames, config["keyboard_dim"])), "mouse": torch.zeros((num_frames, 2)) } grid_sizes = torch.tensor([150, 44, 80]) - + output_dir = os.path.abspath("outputs/matrixgame2") os.makedirs(output_dir, exist_ok=True) video_path = os.path.join(output_dir, f"video_{int(time.time())}.mp4") - + generator.reset( prompt="", image_path=config["image_url"], @@ -473,7 +478,7 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load num_inference_steps=50, output_path=video_path, ) - + new_state = { "initialized": True, "current_model": model_name, @@ -484,13 +489,13 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load "mode": config["mode"], "seed": seed_val, } - + return new_state, seed_val, "Block: 0 / 50", None, gr.update(value="Step"), gr.update(interactive=True) - + async def step_game(keyboard_key, mouse_key, model_name, state): if not state.get("initialized"): return state, state.get("seed", 0), "Block: 0 / 50", None, gr.update(), gr.update() - + # total_start_time = time.time() config = VARIANT_CONFIG.get(model_name) generator = generators.get(config["model_path"]) @@ -500,19 +505,19 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load # Parse inputs to tensors action = get_action_tensors(mode, keyboard_key, mouse_key) keyboard_cond, mouse_cond = expand_action_to_frames(action, frames_per_block) - + # run step async # inference_start_time = time.time() frames, block_future = await generator.step_async(keyboard_cond, mouse_cond) # inference_time = time.time() - inference_start_time - + # wait for block file to be written block_path = await asyncio.to_thread(block_future.result) if block_future else None state["block_idx"] = generator.block_idx block_str = f"Block: {state['block_idx']} / {state['max_blocks']}" - + # total_time = time.time() - total_start_time - + # Timing breakdown # timing_html = create_timing_display(inference_time, total_time, [], frames_per_block) @@ -520,56 +525,58 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load def stop_game(model_name, state): if not state.get("initialized"): - return {"initialized": False}, 0, "Block: 0 / 50", None, gr.update(value="Start"), gr.update(interactive=False) - + return { + "initialized": False + }, 0, "Block: 0 / 50", None, gr.update(value="Start"), gr.update(interactive=False) + config = VARIANT_CONFIG.get(model_name) generator = generators.get(config["model_path"]) - + final_path = state.get("video_path") generator.finalize(final_path) - - return {"initialized": False}, state.get("seed", 0), "Block: 0 / 50", final_path, gr.update(value="Start"), gr.update(interactive=False) - + + return { + "initialized": False + }, state.get("seed", 0), "Block: 0 / 50", final_path, gr.update(value="Start"), gr.update(interactive=False) + async def handle_action(keyboard_key, mouse_key, model_name, seed_val, randomize, state): if not state.get("initialized"): return start_game(model_name, seed_val, randomize, state) else: return await step_game(keyboard_key, mouse_key, model_name, state) - - action_btn.click( - fn=handle_action, - inputs=[keyboard_action, mouse_action, model_selection, seed, randomize_seed, game_state], - outputs=[game_state, seed_output, block_counter, video_output, action_btn, stop_btn] - ) - - stop_btn.click( - fn=stop_game, - inputs=[model_selection, game_state], - outputs=[game_state, seed_output, block_counter, video_output, action_btn, stop_btn] - ) - + + action_btn.click(fn=handle_action, + inputs=[keyboard_action, mouse_action, model_selection, seed, randomize_seed, game_state], + outputs=[game_state, seed_output, block_counter, video_output, action_btn, stop_btn]) + + stop_btn.click(fn=stop_game, + inputs=[model_selection, game_state], + outputs=[game_state, seed_output, block_counter, video_output, action_btn, stop_btn]) + gr.HTML("""

Note that this demo is meant to showcase Matrix Game's quality and that under a large number of requests, generation speed may be affected.

""") - + return demo def main(): parser = argparse.ArgumentParser(description="Matrix Game Gradio Demo") - parser.add_argument("--model", type=str, default="Matrix-Game-2.0-Base", + parser.add_argument("--model", + type=str, + default="Matrix-Game-2.0-Base", choices=list(VARIANT_CONFIG.keys()), help="Model variant to load") parser.add_argument("--host", type=str, default="0.0.0.0") parser.add_argument("--port", type=int, default=7860) args = parser.parse_args() - + # Load the selected model config = VARIANT_CONFIG[args.model] model_path = config["model_path"] - + print(f"Loading model: {model_path}") setup_model_environment(model_path) generator = StreamingVideoGenerator.from_pretrained( @@ -581,43 +588,39 @@ def main(): text_encoder_cpu_offload=True, pin_cpu_memory=True, ) - + generators = {model_path: generator} - + demo = create_gradio_interface(generators, args.model) - + print(f"Starting Gradio at http://{args.host}:{args.port}") - + # FastAPI Wrapper app = FastAPI() - + @app.get("/logo.png") def get_logo(): - return FileResponse( - "assets/full.svg", - media_type="image/svg+xml", - headers={ - "Cache-Control": "public, max-age=3600", - "Access-Control-Allow-Origin": "*" - } - ) + return FileResponse("assets/full.svg", + media_type="image/svg+xml", + headers={ + "Cache-Control": "public, max-age=3600", + "Access-Control-Allow-Origin": "*" + }) @app.get("/favicon.ico") def get_favicon(): favicon_path = "assets/icon-simple.svg" - + if os.path.exists(favicon_path): - return FileResponse( - favicon_path, - media_type="image/svg+xml", - headers={ - "Cache-Control": "public, max-age=3600", - "Access-Control-Allow-Origin": "*" - } - ) + return FileResponse(favicon_path, + media_type="image/svg+xml", + headers={ + "Cache-Control": "public, max-age=3600", + "Access-Control-Allow-Origin": "*" + }) else: raise HTTPException(status_code=404, detail="Favicon not found") - + @app.get("/", response_class=HTMLResponse) def index(request: Request): base_url = str(request.base_url).rstrip('/') @@ -670,13 +673,12 @@ def main(): """ - app = gr.mount_gradio_app( - app, - demo, - path="/gradio", - allowed_paths=[os.path.abspath("outputs"), os.path.abspath("fastvideo-logos")] - ) - + app = gr.mount_gradio_app(app, + demo, + path="/gradio", + allowed_paths=[os.path.abspath("outputs"), + os.path.abspath("fastvideo-logos")]) + uvicorn.run(app, host=args.host, port=args.port) diff --git a/examples/inference/gradio/serving/gradio_frontend.py b/examples/inference/gradio/serving/gradio_frontend.py index bb3f89a96..23f78f853 100644 --- a/examples/inference/gradio/serving/gradio_frontend.py +++ b/examples/inference/gradio/serving/gradio_frontend.py @@ -11,7 +11,6 @@ import gradio as gr from fastvideo.api.sampling_param import SamplingParam - MODEL_PATH_MAPPING = { "FastWan2.1-T2V-1.3B": "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", "FastWan2.2-TI2V-5B-FullAttn": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", @@ -20,32 +19,33 @@ MODEL_PATH_MAPPING = { class RayServeClient: + def __init__(self, backend_url: str): self.backend_url = backend_url self.session = requests.Session() - + def check_health(self) -> bool: try: response = self.session.get(f"{self.backend_url}/health", timeout=5) return response.status_code == 200 except requests.exceptions.RequestException: return False - + def generate_video(self, request_data: dict) -> dict: start_time = time.time() - + try: headers = {"Content-Type": "application/json"} - + response = self.session.post( f"{self.backend_url}/generate_video", json=request_data, headers=headers, timeout=900 # 15 minutes timeout for longer video generation ) - + round_trip_time = time.time() - start_time - + if response.status_code == 200: result = response.json() backend_total = result.get("total_time", 0) @@ -54,7 +54,7 @@ class RayServeClient: return result else: return {"success": False, "error_message": f"HTTP {response.status_code}: {response.text}"} - + except requests.exceptions.RequestException as e: return {"success": False, "error_message": f"Request failed: {str(e)}"} @@ -62,24 +62,24 @@ class RayServeClient: def save_video_from_base64(video_data: str, output_dir: str, prompt: str) -> str: if not video_data: return None - + try: if video_data.startswith('data:video/'): video_data = video_data.split(',')[1] - + video_bytes = base64.b64decode(video_data) - + safe_prompt = prompt[:50].replace(' ', '_').replace('/', '_').replace('\\', '_') video_filename = f"{safe_prompt}.mp4" video_path = os.path.join(output_dir, video_filename) - + os.makedirs(output_dir, exist_ok=True) - + with open(video_path, 'wb') as f: f.write(video_bytes) - + return video_path - + except Exception as e: print(f"Failed to save video: {e}") return None @@ -89,13 +89,13 @@ def encode_image_to_base64(image_path: str) -> str: """Encode an image file to base64 string.""" if not image_path or not os.path.exists(image_path): return None - + try: with open(image_path, 'rb') as f: image_bytes = f.read() - + image_base64 = base64.b64encode(image_bytes).decode('utf-8') - + # Determine image type from extension ext = os.path.splitext(image_path)[1].lower() mime_types = { @@ -106,9 +106,9 @@ def encode_image_to_base64(image_path: str) -> str: '.webp': 'image/webp', } mime_type = mime_types.get(ext, 'image/jpeg') - + return f"data:{mime_type};base64,{image_base64}" - + except Exception as e: print(f"Failed to encode image: {e}") return None @@ -116,7 +116,7 @@ def encode_image_to_base64(image_path: str) -> str: # def create_timing_display(inference_time, encoding_time, network_time, total_time, stage_execution_times, num_frames): # dit_denoising_time = f"{stage_execution_times[5]:.2f}s" if len(stage_execution_times) > 5 else "N/A" -# +# # timing_html = f""" #
#

⏱️ Timing Breakdown

@@ -147,7 +147,7 @@ def encode_image_to_base64(image_path: str) -> str: #
{total_time:.2f}s
#
# """ -# +# # if inference_time > 0: # fps = num_frames / inference_time # timing_html += f""" @@ -155,14 +155,15 @@ def encode_image_to_base64(image_path: str) -> str: # Generation Speed: # {fps:.1f} frames/second # """ -# +# # return timing_html + "" def load_example_prompts(): + def contains_chinese(text): return any('\u4e00' <= char <= '\u9fff' for char in text) - + def load_from_file(filepath): prompts, labels = [], [] try: @@ -176,28 +177,29 @@ def load_example_prompts(): except Exception as e: print(f"Warning: Could not read {filepath}: {e}") return prompts, labels - + # Load prompts from prompts.txt examples, example_labels = load_from_file("examples/inference/gradio/serving/prompts.txt") - + if not examples: - examples = ["A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background."] + examples = [ + "A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background." + ] example_labels = ["Crowded rooftop bar at night"] - + # Load image mappings from JSON file prompt_to_image = {} # Try to find the JSON file relative to project root possible_json_paths = [ Path("assets/prompts/mixkit_i2v.jsonl"), - Path(__file__).resolve().parents[4] / "assets" / "prompts" / - "mixkit_i2v.jsonl", + Path(__file__).resolve().parents[4] / "assets" / "prompts" / "mixkit_i2v.jsonl", ] json_path = None for path in possible_json_paths: if path.exists(): json_path = path break - + if json_path and json_path.exists(): try: with open(json_path, "r", encoding='utf-8') as f: @@ -214,7 +216,7 @@ def load_example_prompts(): prompt_to_image[prompt_text] = str(full_image_path.absolute()) except Exception as e: print(f"Warning: Could not load image mappings from {json_path}: {e}") - + # Create image paths list matching the prompts example_images = [] for prompt in examples: @@ -229,40 +231,38 @@ def load_example_prompts(): image_path = img_path break example_images.append(image_path if image_path and os.path.exists(image_path) else None) - + return examples, example_labels, example_images def create_gradio_interface(backend_url: str, default_params: dict[str, SamplingParam]): - + client = RayServeClient(backend_url) - + def is_i2v_model(model_name: str) -> bool: """Check if the model is an I2V model.""" return "I2V" in model_name - - def generate_video( - prompt, negative_prompt, use_negative_prompt, guidance_scale, - num_frames, height, width, model_selection, input_image, progress - ): + + def generate_video(prompt, negative_prompt, use_negative_prompt, guidance_scale, num_frames, height, width, + model_selection, input_image, progress): # Use default seed value (randomize_seed disabled) seed = 1000 randomize_seed = False if not client.check_health(): return None, f"Backend is not available. Please check if Ray Serve is running at {backend_url}", "" - + # Check if I2V model requires an image if is_i2v_model(model_selection) and not input_image: return None, "I2V models require an input image. Please upload an image.", "" - + # Validate dimensions max_pixels = 720 * 1280 if height * width > max_pixels: return None, f"Video dimensions too large. Maximum: 720x1280 pixels", "" - + if progress: progress(0.1, desc="Checking backend health...") - + # Encode image if provided image_data = None if input_image: @@ -271,7 +271,7 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling image_data = encode_image_to_base64(input_image) if not image_data: return None, "Failed to encode input image", "" - + request_data = { "prompt": prompt, "negative_prompt": negative_prompt, @@ -286,15 +286,15 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling "image_data": image_data, "model_path": MODEL_PATH_MAPPING.get(model_selection, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers") } - + if progress: progress(0.4, desc="Generating video...") - + response = client.generate_video(request_data) - + if progress: progress(0.8, desc="Processing response...") - + if response.get("success", False): video_data = response.get("video_data", "") used_seed = response.get("seed", seed) @@ -303,21 +303,21 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling # total_time = response.get("total_time", 0.0) # network_time = response.get("network_time", 0.0) # stage_execution_times = response.get("stage_execution_times", []) - + # timing_details = create_timing_display( - # inference_time, encoding_time, network_time, total_time, + # inference_time, encoding_time, network_time, total_time, # stage_execution_times, num_frames # ) - + if video_data: if progress: progress(0.9, desc="Saving video...") - + video_path = save_video_from_base64(video_data, "outputs", prompt) - + if progress: progress(1.0, desc="Generation complete!") - + if video_path and os.path.exists(video_path): return video_path, used_seed, "" else: @@ -327,9 +327,9 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling else: error_msg = response.get("error_message", "Unknown error occurred") return None, f"Generation failed: {error_msg}", "" - + examples, example_labels, example_images = load_example_prompts() - + theme = gr.themes.Base().set( button_primary_background_fill="#2563eb", button_primary_background_fill_hover="#1d4ed8", @@ -337,7 +337,7 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling slider_color="#2563eb", checkbox_background_color_selected="#2563eb", ) - + def get_default_values(model_name): # model_path = MODEL_PATH_MAPPING.get(model_name) # if model_path and model_path in default_params: @@ -348,24 +348,24 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling # 'num_frames': params.num_frames, # 'guidance_scale': params.guidance_scale, # } - + return { 'height': 480, 'width': 832, 'num_frames': 73, } - + # Get available models based on what's loaded available_models = [] for model_name, model_path in MODEL_PATH_MAPPING.items(): if model_path in default_params: available_models.append(model_name) - + # Select first available model as default default_model = available_models[0] if available_models else "FastWan2.1-T2V-1.3B" initial_values = get_default_values(default_model) initial_show_image = is_i2v_model(default_model) - + with gr.Blocks(title="CausalWan", theme=theme) as demo: gr.Image("assets/logos/logo.svg", show_label=False, container=False, height=80) gr.HTML(""" @@ -374,7 +374,7 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling

Code | Blog | Docs

""") - + with gr.Accordion("🎥 What Is FastVideo?", open=False): gr.HTML("""
@@ -383,24 +383,20 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling

""") - - with gr.Row(): - model_selection = gr.Dropdown( - choices=available_models, - value=default_model, - label="Select Model", - interactive=True - ) with gr.Row(): - example_dropdown = gr.Dropdown( - choices=example_labels, - label="Example Prompts", - value=None, - interactive=True, - allow_custom_value=False - ) - + model_selection = gr.Dropdown(choices=available_models, + value=default_model, + label="Select Model", + interactive=True) + + with gr.Row(): + example_dropdown = gr.Dropdown(choices=example_labels, + label="Example Prompts", + value=None, + interactive=True, + allow_custom_value=False) + with gr.Row(): with gr.Column(scale=6): prompt = gr.Text( @@ -414,7 +410,7 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling ) with gr.Column(scale=1, min_width=120, elem_classes="center-button"): run_button = gr.Button("Run", variant="primary", size="lg") - + with gr.Row(): with gr.Column(): error_output = gr.Text(label="Error", visible=False) @@ -430,40 +426,31 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling type="filepath", height=400, ) - + with gr.Tab("Advanced Options"): with gr.Group(): with gr.Row(): - height = gr.Number( - label="Height", - value=initial_values['height'], - interactive=False, - container=True - ) - width = gr.Number( - label="Width", - value=initial_values['width'], - interactive=False, - container=True - ) - + height = gr.Number(label="Height", + value=initial_values['height'], + interactive=False, + container=True) + width = gr.Number(label="Width", + value=initial_values['width'], + interactive=False, + container=True) + with gr.Row(): - num_frames = gr.Number( - label="Number of Frames", - value=initial_values['num_frames'], - interactive=False, - container=True - ) - guidance_scale = gr.Number( - label="Guidance Scale", - value=1.0, - interactive=False, - container=True - ) - + num_frames = gr.Number(label="Number of Frames", + value=initial_values['num_frames'], + interactive=False, + container=True) + guidance_scale = gr.Number(label="Guidance Scale", + value=1.0, + interactive=False, + container=True) + with gr.Row(): - use_negative_prompt = gr.Checkbox( - label="Use negative prompt", value=False) + use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False) negative_prompt = gr.Text( label="Negative prompt", max_lines=3, @@ -474,16 +461,16 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling # randomize_seed = gr.Checkbox(label="Randomize seed", value=False) seed_output = gr.Number(label="Used Seed", value=1000) - + with gr.Column(scale=1): result = gr.Video( - label="Generated Video", + label="Generated Video", show_label=True, height=500, container=True, autoplay=True, ) - + gr.HTML(""" """) - + def on_example_select(example_label): if example_label and example_label in example_labels: index = example_labels.index(example_label) @@ -516,32 +503,32 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling selected_image = example_images[index] if index < len(example_images) else None return selected_prompt, selected_image return "", None - + example_dropdown.change( fn=on_example_select, inputs=example_dropdown, outputs=[prompt, input_image], ) - + gr.HTML("""

The compute for this demo is generously provided by GMI Cloud. Note that this demo is meant as a preview of our distilled I2V model. Outside of few-step distillation, we have not yet fully optimized it for speed. Stay tuned for updates!

""") - + use_negative_prompt.change( fn=lambda x: gr.update(visible=x), inputs=use_negative_prompt, outputs=negative_prompt, ) - + def on_model_selection_change(selected_model): if not selected_model: selected_model = "FastWan2.1-T2V-1.3B" - + model_path = MODEL_PATH_MAPPING.get(selected_model) show_image_input = is_i2v_model(selected_model) - + if model_path and model_path in default_params: params = default_params[model_path] return ( @@ -551,7 +538,7 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling gr.update(value=params.guidance_scale), gr.update(visible=show_image_input), ) - + return ( gr.update(value=448), gr.update(value=832), @@ -559,34 +546,33 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling gr.update(value=3.0), gr.update(visible=show_image_input), ) - + model_selection.change( fn=on_model_selection_change, inputs=model_selection, outputs=[height, width, num_frames, guidance_scale, image_tab], ) - + def handle_generation(*args, progress=None, request: gr.Request = None): model_selection, prompt, negative_prompt, use_negative_prompt, guidance_scale, num_frames, height, width, input_image = args - - result_path, seed_or_error, _ = generate_video( - prompt, negative_prompt, use_negative_prompt, guidance_scale, - num_frames, height, width, model_selection, input_image, progress - ) - + + result_path, seed_or_error, _ = generate_video(prompt, negative_prompt, use_negative_prompt, guidance_scale, + num_frames, height, width, model_selection, input_image, + progress) + if result_path and os.path.exists(result_path): return ( - result_path, - seed_or_error, + result_path, + seed_or_error, gr.update(visible=False), ) else: return ( - None, - seed_or_error, + None, + seed_or_error, gr.update(visible=True, value=seed_or_error), ) - + run_button.click( fn=handle_generation, inputs=[ @@ -604,80 +590,75 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling outputs=[result, seed_output, error_output], # timing_display removed concurrency_limit=20, ) - + return demo def main(): parser = argparse.ArgumentParser(description="FastVideo Gradio Frontend") - parser.add_argument("--backend_url", type=str, default="http://localhost:8000", - help="URL of the Ray Serve backend") - parser.add_argument("--t2v_model_paths", type=str, + parser.add_argument("--backend_url", type=str, default="http://localhost:8000", help="URL of the Ray Serve backend") + parser.add_argument("--t2v_model_paths", + type=str, default="", help="Comma separated list of paths to the T2V model(s)") - parser.add_argument("--i2v_model_paths", type=str, + parser.add_argument("--i2v_model_paths", + type=str, default="FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers", help="Comma separated list of paths to the I2V model(s)") - parser.add_argument("--host", type=str, default="0.0.0.0", - help="Host to bind to") - parser.add_argument("--port", type=int, default=7860, - help="Port to bind to") - + parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind to") + parser.add_argument("--port", type=int, default=7860, help="Port to bind to") + args = parser.parse_args() - + default_params = {} - + # Load T2V model params t2v_paths = [p.strip() for p in args.t2v_model_paths.split(",") if p.strip()] for model_path in t2v_paths: default_params[model_path] = SamplingParam.from_pretrained(model_path) - + # Load I2V model params i2v_paths = [p.strip() for p in args.i2v_model_paths.split(",") if p.strip()] for model_path in i2v_paths: default_params[model_path] = SamplingParam.from_pretrained(model_path) - + demo = create_gradio_interface(args.backend_url, default_params) - + print(f"Starting Gradio frontend at http://{args.host}:{args.port}") print(f"Backend URL: {args.backend_url}") print(f"T2V Models: {args.t2v_model_paths}") if args.i2v_model_paths: print(f"I2V Models: {args.i2v_model_paths}") - + from fastapi import FastAPI, Request, HTTPException from fastapi.responses import HTMLResponse, FileResponse import uvicorn - + app = FastAPI() - + @app.get("/logo.svg") def get_logo(): - return FileResponse( - "assets/logos/logo.svg", - media_type="image/svg+xml", - headers={ - "Cache-Control": "public, max-age=3600", - "Access-Control-Allow-Origin": "*" - } - ) - + return FileResponse("assets/logos/logo.svg", + media_type="image/svg+xml", + headers={ + "Cache-Control": "public, max-age=3600", + "Access-Control-Allow-Origin": "*" + }) + @app.get("/favicon.ico") def get_favicon(): favicon_path = "assets/logos/icon_simple.svg" - + if os.path.exists(favicon_path): - return FileResponse( - favicon_path, - media_type="image/svg+xml", - headers={ - "Cache-Control": "public, max-age=3600", - "Access-Control-Allow-Origin": "*" - } - ) + return FileResponse(favicon_path, + media_type="image/svg+xml", + headers={ + "Cache-Control": "public, max-age=3600", + "Access-Control-Allow-Origin": "*" + }) else: raise HTTPException(status_code=404, detail="Favicon not found") - + @app.get("/", response_class=HTMLResponse) def index(request: Request): base_url = str(request.base_url).rstrip('/') @@ -729,23 +710,21 @@ def main(): """ - - app = gr.mount_gradio_app( - app, - demo, - path="/gradio", - allowed_paths=[ - os.path.abspath("outputs"), - os.path.abspath("fastvideo-logos"), - os.path.abspath("assets/prompts"), - os.path.abspath("assets/images"), - os.path.abspath(tempfile.gettempdir()), - os.path.abspath(os.path.join(tempfile.gettempdir(), "gradio")), - ] - ) - + + app = gr.mount_gradio_app(app, + demo, + path="/gradio", + allowed_paths=[ + os.path.abspath("outputs"), + os.path.abspath("fastvideo-logos"), + os.path.abspath("assets/prompts"), + os.path.abspath("assets/images"), + os.path.abspath(tempfile.gettempdir()), + os.path.abspath(os.path.join(tempfile.gettempdir(), "gradio")), + ]) + uvicorn.run(app, host=args.host, port=args.port) if __name__ == "__main__": - main() + main() diff --git a/examples/inference/gradio/serving/ray_serve_backend.py b/examples/inference/gradio/serving/ray_serve_backend.py index 2a783207b..4be1518b4 100644 --- a/examples/inference/gradio/serving/ray_serve_backend.py +++ b/examples/inference/gradio/serving/ray_serve_backend.py @@ -24,7 +24,7 @@ NUM_GPUS = 16 DEFAULT_FPS = 16 SEED_RANGE_MAX = 1_000_000 SUPPORTED_MODELS = [ - "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", + "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", "FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers", ] @@ -86,15 +86,15 @@ class VideoGenerationResponse(BaseModel): def encode_video_to_base64(frames: List[np.ndarray], fps: int = DEFAULT_FPS) -> str: if not frames: return "" - + try: buffer = io.BytesIO() imageio.mimsave(buffer, frames, fps=fps, format="mp4") buffer.seek(0) - + video_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8') return f"data:video/mp4;base64,{video_base64}" - + except Exception as e: print(f"Warning: Failed to encode video: {e}") return "" @@ -104,23 +104,23 @@ def save_image_from_base64(image_data: str, output_dir: str) -> Optional[str]: """Save base64 image data to a temporary file and return the path.""" if not image_data: return None - + try: # Remove data URL prefix if present if image_data.startswith('data:image/'): image_data = image_data.split(',')[1] - + image_bytes = base64.b64decode(image_data) - + # Save to temporary file os.makedirs(output_dir, exist_ok=True) temp_image_path = os.path.join(output_dir, f"temp_input_{int(time.time() * 1000)}.png") - + with open(temp_image_path, 'wb') as f: f.write(image_bytes) - + return temp_image_path - + except Exception as e: print(f"Warning: Failed to save image: {e}") return None @@ -138,30 +138,29 @@ def setup_model_environment(model_path: str) -> None: def process_generation_result(result: Any) -> tuple[List[np.ndarray], float, List[str], List[float]]: frames = result if isinstance(result, list) else result.get("frames", []) generation_time = result.get("generation_time", 0.0) if isinstance(result, dict) else 0.0 - + logging_info = result.get("logging_info", None) if logging_info: stage_names = logging_info.get_execution_order() stage_execution_times = [ - logging_info.get_stage_info(stage_name).get("execution_time", 0.0) - for stage_name in stage_names + logging_info.get_stage_info(stage_name).get("execution_time", 0.0) for stage_name in stage_names ] else: stage_names = [] stage_execution_times = [] - + return frames, generation_time, stage_names, stage_execution_times def prepare_sampling_params(video_request: VideoGenerationRequest, default_params: Any) -> Any: params = deepcopy(default_params) params.prompt = video_request.prompt - + if video_request.use_negative_prompt: params.negative_prompt = video_request.negative_prompt - params.seed = (video_request.seed if not video_request.randomize_seed - else torch.randint(0, SEED_RANGE_MAX, (1,)).item()) + params.seed = (video_request.seed if not video_request.randomize_seed else torch.randint(0, SEED_RANGE_MAX, + (1, )).item()) params.randomize_seed = video_request.randomize_seed params.guidance_scale = video_request.guidance_scale params.num_frames = video_request.num_frames @@ -169,11 +168,12 @@ def prepare_sampling_params(video_request: VideoGenerationRequest, default_param params.width = video_request.width params.save_video = False params.return_frames = True - + return params class BaseModelDeployment: + def __init__(self, model_path: str, output_path: str = "outputs"): self.model_path = model_path self.output_path = output_path @@ -193,7 +193,7 @@ class BaseModelDeployment: num_gpus=1, use_fsdp_inference=True, text_encoder_cpu_offload=config["text_encoder_cpu_offload"], - dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125], # TODO: hardocde for I2V + dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125], # TODO: hardocde for I2V dit_precision="fp32", # TODO: hardocde for I2V dit_cpu_offload=config["dit_cpu_offload"], vae_cpu_offload=config["vae_cpu_offload"], @@ -208,7 +208,7 @@ class BaseModelDeployment: def generate_video(self, video_request: VideoGenerationRequest) -> VideoGenerationResponse: total_start_time = time.time() - + params = prepare_sampling_params(video_request, self.default_params) # Save image if provided (for I2V) @@ -238,9 +238,9 @@ class BaseModelDeployment: encoding_start_time = time.time() video_data = encode_video_to_base64(frames, fps=DEFAULT_FPS) encoding_time = time.time() - encoding_start_time - + total_time = time.time() - total_start_time - + # Clean up temporary image file if image_path and os.path.exists(image_path): try: @@ -261,20 +261,18 @@ class BaseModelDeployment: ) -@serve.deployment( - ray_actor_options={"num_cpus": 2, "num_gpus": 1, "runtime_env": {"conda": "demo-fv"}}, -) +@serve.deployment(ray_actor_options={"num_cpus": 2, "num_gpus": 1, "runtime_env": {"conda": "demo-fv"}}, ) class T2VModelDeployment(BaseModelDeployment): + def __init__(self, t2v_model_path: str, output_path: str = "outputs"): super().__init__(t2v_model_path, output_path) self._initialize_generator(MODEL_CONFIGS["1.3B"]) print("✅ T2V model initialized successfully") -@serve.deployment( - ray_actor_options={"num_cpus": 16, "num_gpus": 1, "runtime_env": {"conda": "demo-fv"}}, -) +@serve.deployment(ray_actor_options={"num_cpus": 16, "num_gpus": 1, "runtime_env": {"conda": "demo-fv"}}, ) class T2V14BModelDeployment(BaseModelDeployment): + def __init__(self, t2v_14b_model_path: str, output_path: str = "outputs"): super().__init__(t2v_14b_model_path, output_path) # Override environment for 14B model @@ -283,10 +281,9 @@ class T2V14BModelDeployment(BaseModelDeployment): print("✅ T2V 14B model initialized successfully") -@serve.deployment( - ray_actor_options={"num_cpus": 15, "num_gpus": 1, "runtime_env": {"conda": "demo-fv"}}, -) +@serve.deployment(ray_actor_options={"num_cpus": 15, "num_gpus": 1, "runtime_env": {"conda": "demo-fv"}}, ) class I2VModelDeployment(BaseModelDeployment): + def __init__(self, i2v_model_path: str, output_path: str = "outputs"): super().__init__(i2v_model_path, output_path) # Override environment for I2V model @@ -305,46 +302,54 @@ app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler) @serve.ingress(app) class FastVideoAPI: - def __init__(self, t2v_deployments: Dict[str, DeploymentHandle], i2v_deployments: Dict[str, DeploymentHandle] = None): + def __init__(self, + t2v_deployments: Dict[str, DeploymentHandle], + i2v_deployments: Dict[str, DeploymentHandle] = None): self.t2v_deployments = t2v_deployments self.i2v_deployments = i2v_deployments or {} self.all_deployments = {**self.t2v_deployments, **self.i2v_deployments} - + # Initialize Prometheus metrics self.request_count = Counter('fastvideo_requests_total', 'Total FastVideo requests', ['model_type', 'status']) - self.request_duration = Histogram('fastvideo_request_duration_seconds', 'FastVideo request duration', ['model_type']) - self.video_generation_time = Histogram('fastvideo_video_generation_seconds', 'Video generation time', ['model_type']) - + self.request_duration = Histogram('fastvideo_request_duration_seconds', 'FastVideo request duration', + ['model_type']) + self.video_generation_time = Histogram('fastvideo_video_generation_seconds', 'Video generation time', + ['model_type']) + def _get_model_name(self, model_path: Optional[str]) -> str: return model_path.split('/')[-1] if model_path else "unknown" - - def _record_metrics(self, model_name: str, status: str, duration: float, response: Optional[VideoGenerationResponse] = None) -> None: + + def _record_metrics(self, + model_name: str, + status: str, + duration: float, + response: Optional[VideoGenerationResponse] = None) -> None: self.request_count.labels(model_type=model_name, status=status).inc() self.request_duration.labels(model_type=model_name).observe(duration) - + if response and hasattr(response, 'generation_time') and response.generation_time: self.video_generation_time.labels(model_type=model_name).observe(response.generation_time) - + @app.post("/generate_video", response_model=VideoGenerationResponse) @limiter.limit("10/minute") async def generate_video(self, request: Request, video_request: VideoGenerationRequest) -> VideoGenerationResponse: """Route the request to the appropriate model deployment based on model_path.""" start_time = time.time() model_name = self._get_model_name(video_request.model_path) - + try: if video_request.model_path not in self.all_deployments: raise ValueError(f"Model {video_request.model_path} not found") - + response_ref = self.all_deployments[video_request.model_path].generate_video.remote(video_request) response = await response_ref - + self._record_metrics(model_name, "success", time.time() - start_time, response) return response except Exception as e: self._record_metrics(model_name, "error", time.time() - start_time) - + return VideoGenerationResponse( video_data=None, seed=video_request.seed, @@ -355,12 +360,12 @@ class FastVideoAPI: encoding_time=0, total_time=0, ) - + @app.get("/health") @limiter.limit("10/minute") async def health_check(self, request: Request) -> Dict[str, str]: return {"status": "healthy"} - + @app.get("/metrics") async def metrics(self) -> Response: return Response(generate_latest(), media_type="text/plain") @@ -370,7 +375,7 @@ def validate_configuration(model_paths: List[str], replicas: List[int]) -> None: assert len(model_paths) > 0, "At least one model must be specified" assert len(model_paths) == len(replicas), "Number of models and replicas must match" assert sum(replicas) <= NUM_GPUS, f"Total replicas ({sum(replicas)}) must be <= {NUM_GPUS}" - + for model, replica_count in zip(model_paths, replicas): assert model in SUPPORTED_MODELS, f"Model {model} not supported. Supported models: {SUPPORTED_MODELS}" assert replica_count > 0, f"Replicas must be greater than 0" @@ -392,11 +397,11 @@ def start_ray_serve( # Parse T2V models t2v_paths = [p.strip() for p in t2v_model_paths.split(",") if p.strip()] t2v_reps = [int(r.strip()) for r in t2v_model_replicas.split(",") if r.strip()] if t2v_model_replicas else [] - + # Parse I2V models i2v_paths = [p.strip() for p in i2v_model_paths.split(",") if p.strip()] i2v_reps = [int(r.strip()) for r in i2v_model_replicas.split(",") if r.strip()] if i2v_model_replicas else [] - + # Validate configurations all_paths = t2v_paths + i2v_paths all_replicas = t2v_reps + i2v_reps @@ -433,7 +438,7 @@ def setup_signal_handlers() -> None: if __name__ == "__main__": import argparse - + parser = argparse.ArgumentParser(description="FastVideo Ray Serve Backend") parser.add_argument("--t2v_model_paths", type=str, @@ -451,31 +456,24 @@ if __name__ == "__main__": type=str, default="1", help="Comma separated list of number of replicas for the I2V model(s)") - parser.add_argument("--output_path", - type=str, - default="outputs", - help="Path to save generated videos") - parser.add_argument("--host", - type=str, - default="0.0.0.0", - help="Host to bind to") - parser.add_argument("--port", - type=int, - default=8000, - help="Port to bind to") - + parser.add_argument("--output_path", type=str, default="outputs", help="Path to save generated videos") + parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind to") + parser.add_argument("--port", type=int, default=8000, help="Port to bind to") + args = parser.parse_args() # Parse and validate all models t2v_paths = [p.strip() for p in args.t2v_model_paths.split(",") if p.strip()] - t2v_reps = [int(r.strip()) for r in args.t2v_model_replicas.split(",") if r.strip()] if args.t2v_model_replicas else [] + t2v_reps = [int(r.strip()) for r in args.t2v_model_replicas.split(",") + if r.strip()] if args.t2v_model_replicas else [] i2v_paths = [p.strip() for p in args.i2v_model_paths.split(",") if p.strip()] - i2v_reps = [int(r.strip()) for r in args.i2v_model_replicas.split(",") if r.strip()] if args.i2v_model_replicas else [] - + i2v_reps = [int(r.strip()) for r in args.i2v_model_replicas.split(",") + if r.strip()] if args.i2v_model_replicas else [] + all_paths = t2v_paths + i2v_paths all_replicas = t2v_reps + i2v_reps validate_configuration(all_paths, all_replicas) - + start_ray_serve( t2v_model_paths=args.t2v_model_paths, t2v_model_replicas=args.t2v_model_replicas, @@ -489,4 +487,4 @@ if __name__ == "__main__": setup_signal_handlers() print("✅ FastVideo backend is running. Press Ctrl-C to stop.") while True: - time.sleep(3600) + time.sleep(3600) diff --git a/examples/inference/gradio/serving/start_ray_serve_app.py b/examples/inference/gradio/serving/start_ray_serve_app.py index 5790b922d..ea799db48 100644 --- a/examples/inference/gradio/serving/start_ray_serve_app.py +++ b/examples/inference/gradio/serving/start_ray_serve_app.py @@ -14,7 +14,6 @@ import requests from pathlib import Path from typing import Dict, Any, Optional - DEFAULT_BACKEND_HOST = "0.0.0.0" DEFAULT_BACKEND_PORT = 8000 DEFAULT_FRONTEND_HOST = "0.0.0.0" @@ -36,16 +35,16 @@ sys.path.insert(0, str(PROJECT_ROOT)) class ServiceManager: - + def __init__(self, args: argparse.Namespace): self.args = args self.backend_process: Optional[subprocess.Popen] = None self.frontend_process: Optional[subprocess.Popen] = None self.backend_url = f"http://{args.backend_host}:{args.backend_port}" - + def check_backend_health(self, max_retries: int = HEALTH_CHECK_MAX_RETRIES) -> bool: health_url = f"{self.backend_url}/health" - + for attempt in range(max_retries): try: response = requests.get(health_url, timeout=HEALTH_CHECK_TIMEOUT) @@ -54,45 +53,44 @@ class ServiceManager: return True except requests.exceptions.RequestException: pass - + if attempt < max_retries - 1: print(f"⏳ Waiting for backend to start... ({attempt + 1}/{max_retries})") time.sleep(HEALTH_CHECK_INTERVAL) - + print(f"❌ Backend failed to start within {max_retries * HEALTH_CHECK_INTERVAL} seconds") return False - + def _create_monitor_thread(self, process: subprocess.Popen, service_name: str) -> threading.Thread: + def monitor(): if process.stdout: for line in process.stdout: print(f"[{service_name}] {line.rstrip()}") - + thread = threading.Thread(target=monitor, daemon=True) thread.start() return thread - + def _start_service(self, script_name: str, args_dict: Dict[str, Any], service_name: str) -> subprocess.Popen: script_path = Path(__file__).parent / script_name - + cmd = [sys.executable, str(script_path)] for key, value in args_dict.items(): cmd.extend([f"--{key}", str(value)]) - + print(f"🚀 Starting {service_name}...") print(f"Command: {' '.join(cmd)}") - - process = subprocess.Popen( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - universal_newlines=True, - bufsize=1 - ) - + + process = subprocess.Popen(cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + bufsize=1) + self._create_monitor_thread(process, service_name.upper()) return process - + def start_backend(self) -> subprocess.Popen: backend_args = { "t2v_model_paths": self.args.t2v_model_paths, @@ -101,16 +99,16 @@ class ServiceManager: "host": self.args.backend_host, "port": self.args.backend_port } - + # Add I2V parameters if provided if self.args.i2v_model_paths: backend_args["i2v_model_paths"] = self.args.i2v_model_paths if self.args.i2v_model_replicas: backend_args["i2v_model_replicas"] = self.args.i2v_model_replicas - + self.backend_process = self._start_service("ray_serve_backend.py", backend_args, "backend") return self.backend_process - + def start_frontend(self) -> subprocess.Popen: frontend_args = { "backend_url": self.backend_url, @@ -118,26 +116,26 @@ class ServiceManager: "host": self.args.frontend_host, "port": self.args.frontend_port } - + # Add I2V parameters if provided if self.args.i2v_model_paths: frontend_args["i2v_model_paths"] = self.args.i2v_model_paths - + self.frontend_process = self._start_service("gradio_frontend.py", frontend_args, "frontend") return self.frontend_process - + def shutdown_services(self) -> None: print("\n🛑 Shutting down services...") - + processes = [] if self.frontend_process: self.frontend_process.terminate() processes.append(("frontend", self.frontend_process)) - + if self.backend_process: self.backend_process.terminate() processes.append(("backend", self.backend_process)) - + for name, process in processes: try: process.wait(timeout=PROCESS_SHUTDOWN_TIMEOUT) @@ -145,37 +143,38 @@ class ServiceManager: except subprocess.TimeoutExpired: print(f"⚠️ Force killing {name} process...") process.kill() - + print("✅ All services stopped") - + def monitor_processes(self) -> None: if not self.backend_process or not self.frontend_process: print("❌ Processes not properly initialized") return - + try: while True: if self.frontend_process.poll() is not None: print("❌ Frontend process died unexpectedly") break - + if self.backend_process.poll() is not None: print("❌ Backend process died unexpectedly") break - + time.sleep(PROCESS_MONITOR_INTERVAL) - + except KeyboardInterrupt: pass - + self.shutdown_services() def setup_signal_handlers(service_manager: ServiceManager) -> None: + def signal_handler(signum: int, frame: Any) -> None: service_manager.shutdown_services() sys.exit(0) - + signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) @@ -196,80 +195,63 @@ def print_startup_info(args: argparse.Namespace) -> None: def parse_arguments() -> argparse.Namespace: parser = argparse.ArgumentParser(description="FastVideo Ray Serve App") - - parser.add_argument("--t2v_model_paths", + + parser.add_argument("--t2v_model_paths", type=str, default=DEFAULT_T2V_MODELS, help="Comma separated list of paths to the T2V model(s)") - parser.add_argument("--t2v_model_replicas", - type=str, + parser.add_argument("--t2v_model_replicas", + type=str, default=DEFAULT_T2V_REPLICAS, help="Comma separated list of number of replicas for the T2V model(s)") - parser.add_argument("--i2v_model_paths", + parser.add_argument("--i2v_model_paths", type=str, default=DEFAULT_I2V_MODELS, help="Comma separated list of paths to the I2V model(s)") - parser.add_argument("--i2v_model_replicas", - type=str, + parser.add_argument("--i2v_model_replicas", + type=str, default=DEFAULT_I2V_REPLICAS, help="Comma separated list of number of replicas for the I2V model(s)") - parser.add_argument("--output_path", - type=str, - default=DEFAULT_OUTPUT_PATH, - help="Path to save generated videos") - - parser.add_argument("--backend_host", - type=str, - default=DEFAULT_BACKEND_HOST, - help="Backend host to bind to") - parser.add_argument("--backend_port", - type=int, - default=DEFAULT_BACKEND_PORT, - help="Backend port to bind to") - - parser.add_argument("--frontend_host", - type=str, - default=DEFAULT_FRONTEND_HOST, - help="Frontend host to bind to") - parser.add_argument("--frontend_port", - type=int, - default=DEFAULT_FRONTEND_PORT, - help="Frontend port to bind to") - - parser.add_argument("--skip_backend_check", - action="store_true", - help="Skip backend health check") - + parser.add_argument("--output_path", type=str, default=DEFAULT_OUTPUT_PATH, help="Path to save generated videos") + + parser.add_argument("--backend_host", type=str, default=DEFAULT_BACKEND_HOST, help="Backend host to bind to") + parser.add_argument("--backend_port", type=int, default=DEFAULT_BACKEND_PORT, help="Backend port to bind to") + + parser.add_argument("--frontend_host", type=str, default=DEFAULT_FRONTEND_HOST, help="Frontend host to bind to") + parser.add_argument("--frontend_port", type=int, default=DEFAULT_FRONTEND_PORT, help="Frontend port to bind to") + + parser.add_argument("--skip_backend_check", action="store_true", help="Skip backend health check") + return parser.parse_args() def main() -> None: args = parse_arguments() - + os.makedirs(args.output_path, exist_ok=True) print_startup_info(args) - + service_manager = ServiceManager(args) setup_signal_handlers(service_manager) - + try: service_manager.start_backend() - + if not args.skip_backend_check: if not service_manager.check_backend_health(): print("❌ Backend failed to start. Terminating...") service_manager.shutdown_services() sys.exit(1) - + service_manager.start_frontend() - + print("\n🎉 Both services are starting up!") print(f"📺 Frontend will be available at: http://{args.frontend_host}:{args.frontend_port}") print(f"🔧 Backend API will be available at: http://{args.backend_host}:{args.backend_port}") print("\nPress Ctrl+C to stop both services...") - + service_manager.monitor_processes() - + except Exception as e: print(f"❌ Unexpected error: {e}") service_manager.shutdown_services() @@ -277,4 +259,4 @@ def main() -> None: if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/inference/lora/wan_lora_inference.py b/examples/inference/lora/wan_lora_inference.py index 79fe4300e..1b42be28b 100644 --- a/examples/inference/lora/wan_lora_inference.py +++ b/examples/inference/lora/wan_lora_inference.py @@ -2,6 +2,8 @@ from fastvideo import VideoGenerator from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "./lora_out" + + def main(): # Initialize VideoGenerator with the Wan model generator = VideoGenerator.from_pretrained( @@ -10,10 +12,9 @@ def main(): dit_cpu_offload=False, vae_cpu_offload=False, text_encoder_cpu_offload=True, - pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" + pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" lora_path="benjamin-paine/steamboat-willie-1.3b", - lora_nickname="steamboat" - ) + lora_nickname="steamboat") kwargs = { "height": 480, "width": 832, @@ -32,19 +33,17 @@ def main(): output_path=OUTPUT_PATH, save_video=True, negative_prompt=negative_prompt, - **kwargs - ) - + **kwargs) + generator.set_lora_adapter(lora_nickname="flat_color", lora_path="motimalu/wan-flat-color-1.3b-v2") prompt = "flat color, no lineart, blending, negative space, artist:[john kafka|ponsuke kaikai|hara id 21|yoneyama mai|fuzichoco], 1girl, sakura miko, pink hair, cowboy shot, white shirt, floral print, off shoulder, outdoors, cherry blossom, tree shade, wariza, looking up, falling petals, half-closed eyes, white sky, clouds, live2d animation, upper body, high quality cinematic video of a woman sitting under a sakura tree. Dreamy and lonely, the camera close-ups on the face of the woman as she turns towards the viewer. The Camera is steady, This is a cowboy shot. The animation is smooth and fluid." negative_prompt = "bad quality video,色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" - video = generator.generate_video( - prompt, - output_path=OUTPUT_PATH, - save_video=True, - negative_prompt=negative_prompt, - **kwargs - ) + video = generator.generate_video(prompt, + output_path=OUTPUT_PATH, + save_video=True, + negative_prompt=negative_prompt, + **kwargs) + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/inference/lora/wan_lora_inference_from_ckpt.py b/examples/inference/lora/wan_lora_inference_from_ckpt.py index ca7032133..99f6bb1c3 100644 --- a/examples/inference/lora/wan_lora_inference_from_ckpt.py +++ b/examples/inference/lora/wan_lora_inference_from_ckpt.py @@ -5,6 +5,8 @@ from fastvideo import VideoGenerator from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "./lora_out" + + def main(): # Initialize VideoGenerator with the Wan model generator = VideoGenerator.from_pretrained( @@ -13,10 +15,9 @@ def main(): dit_cpu_offload=False, vae_cpu_offload=True, text_encoder_cpu_offload=True, - pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" + pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-160/transformer", - lora_nickname="crush_smol" - ) + lora_nickname="crush_smol") generator.unmerge_lora_weights() kwargs = { "height": 480, @@ -29,18 +30,10 @@ def main(): # Generate video with LoRA style prompt = "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press." - video = generator.generate_video( - prompt, - output_path=OUTPUT_PATH, - save_video=True, - **kwargs - ) + video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, **kwargs) prompt = "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press." - video = generator.generate_video( - prompt, - output_path=OUTPUT_PATH, - save_video=True, - **kwargs - ) + video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, **kwargs) + + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/inference/ltx2_3/optimized_nvfp4_t2v.py b/examples/inference/ltx2_3/optimized_nvfp4_t2v.py index 5e326a9eb..ed5161f42 100644 --- a/examples/inference/ltx2_3/optimized_nvfp4_t2v.py +++ b/examples/inference/ltx2_3/optimized_nvfp4_t2v.py @@ -55,18 +55,13 @@ _inductor.coordinate_descent_tuning = True _inductor.coordinate_descent_check_all_directions = True _inductor.epilogue_fusion = False # do not fuse pointwise ops into matmuls -MODEL_ID = os.path.expandvars( - os.path.expanduser( - os.getenv("LTX23_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers") - ) -) +MODEL_ID = os.path.expandvars(os.path.expanduser(os.getenv("LTX23_MODEL_PATH", + "FastVideo/LTX-2.3-Distilled-Diffusers"))) OUTPUT_DIR = Path(os.getenv("LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_nvfp4_t2v")) -DEFAULT_PROMPT = ( - "A fashion model takes a slow step forward and shifts her weight, " - "the soft fabric of her clothing swaying and rippling with the " - "motion, her hair shifting gently, soft even studio lighting on a " - "clean light background, elegant slow-motion runway feel." -) +DEFAULT_PROMPT = ("A fashion model takes a slow step forward and shifts her weight, " + "the soft fabric of her clothing swaying and rippling with the " + "motion, her hair shifting gently, soft even studio lighting on a " + "clean light background, elegant slow-motion runway feel.") PROMPT = os.getenv("LTX23_T2V_PROMPT", DEFAULT_PROMPT) @@ -120,10 +115,12 @@ def main() -> None: common_kwargs = dict( prompt=PROMPT, - negative_prompt="", # distilled is CFG-free; no negative needed - guidance_scale=1.0, # CFG=1 for distilled - height=1280, width=832, # portrait runway aspect - num_frames=121, fps=24, # ~5s clip + negative_prompt="", # distilled is CFG-free; no negative needed + guidance_scale=1.0, # CFG=1 for distilled + height=1280, + width=832, # portrait runway aspect + num_frames=121, + fps=24, # ~5s clip # Single-stage 8-step distilled sampling — the validated preset for # this checkpoint (no two-stage refine; the NVFP4 deploy contract # runs the distilled single-stage recipe). diff --git a/examples/inference/optimizations/FastWan_QAD_TAEHV.py b/examples/inference/optimizations/FastWan_QAD_TAEHV.py index aa4b6840d..f7ada6e44 100644 --- a/examples/inference/optimizations/FastWan_QAD_TAEHV.py +++ b/examples/inference/optimizations/FastWan_QAD_TAEHV.py @@ -46,19 +46,15 @@ OUTPUT_PATH = "video_samples" # guidance 1.0). Loaded on top of the base Wan2.1 pipeline; pass # ``--distilled_model ''`` to run the base weights instead. DEFAULT_DISTILLED_MODEL = "FastVideo/FastWan-QAD-1.3B" -DISTILLED_WEIGHTS_FILE = ( - "generator_inference_transformer/diffusion_pytorch_model.safetensors" -) +DISTILLED_WEIGHTS_FILE = ("generator_inference_transformer/diffusion_pytorch_model.safetensors") # TAEHV checkpoint for Wan2.1. Clone https://github.com/madebyollin/taehv to get # ``taew2_1.pth`` (Wan 2.1 / Wan 2.2-14B / Qwen-Image all use this VAE). DEFAULT_TAEHV_CHECKPOINT = "/root/taehv/taew2_1.pth" -PROMPT = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." -) +PROMPT = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") class TaehvDecoder: @@ -69,8 +65,7 @@ class TaehvDecoder: needed -- unlike the full Wan VAE path. """ - def __init__(self, checkpoint_path: str, device: str = "cuda", - dtype: torch.dtype = torch.float16) -> None: + def __init__(self, checkpoint_path: str, device: str = "cuda", dtype: torch.dtype = torch.float16) -> None: from taehv import TAEHV # pip-installed; no sys.path manipulation self.device = device self.dtype = dtype @@ -89,8 +84,7 @@ class TaehvDecoder: """ # NCTHW -> NTCHW (TAEHV's expected layout), on the TAEHV device/dtype. latents = latents.permute(0, 2, 1, 3, 4).to(self.device, self.dtype) - decoded = self.model.decode_video( - latents, parallel=True, show_progress_bar=False) + decoded = self.model.decode_video(latents, parallel=True, show_progress_bar=False) # decoded: [B, T, 3, H, W] in [0, 1]. Take batch 0, vectorize to uint8. frames = (decoded[0].clamp(0, 1) * 255).to(torch.uint8) return frames.permute(0, 2, 3, 1).cpu().numpy() @@ -127,10 +121,9 @@ def resolve_taehv_checkpoint(path: str) -> str: """Validate the TAEHV checkpoint path, with a helpful error if missing.""" if os.path.exists(path): return path - raise FileNotFoundError( - f"TAEHV checkpoint not found at {path!r}. Clone the weights with:\n" - " git clone https://github.com/madebyollin/taehv\n" - "and pass --taehv_checkpoint /taew2_1.pth") + raise FileNotFoundError(f"TAEHV checkpoint not found at {path!r}. Clone the weights with:\n" + " git clone https://github.com/madebyollin/taehv\n" + "and pass --taehv_checkpoint /taew2_1.pth") def build_generator(args: argparse.Namespace) -> VideoGenerator: @@ -141,7 +134,7 @@ def build_generator(args: argparse.Namespace) -> VideoGenerator: pipeline_config = PipelineConfig.from_pretrained(model_id) pipeline_config.dit_precision = "bf16" pipeline_config.vae_precision = "bf16" - pipeline_config.text_encoder_precisions = ("bf16",) + pipeline_config.text_encoder_precisions = ("bf16", ) if not args.baseline: pipeline_config.dit_config.quant_config = NVFP4QATConfig() @@ -181,30 +174,30 @@ def build_generator(args: argparse.Namespace) -> VideoGenerator: def main() -> None: - parser = argparse.ArgumentParser( - description="FP4 linear Wan2.1-1.3B with TAEHV decoding benchmark") - parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - help="Model path or HuggingFace ID") - parser.add_argument("--baseline", action="store_true", - help="Run dense bf16 instead of FP4 linear") - parser.add_argument("--no-compile", action="store_true", - help="Disable torch.compile (eager)") - parser.add_argument("--taehv", action=argparse.BooleanOptionalAction, + parser = argparse.ArgumentParser(description="FP4 linear Wan2.1-1.3B with TAEHV decoding benchmark") + parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", help="Model path or HuggingFace ID") + parser.add_argument("--baseline", action="store_true", help="Run dense bf16 instead of FP4 linear") + parser.add_argument("--no-compile", action="store_true", help="Disable torch.compile (eager)") + parser.add_argument("--taehv", + action=argparse.BooleanOptionalAction, default=True, help="Decode with TAEHV instead of the full Wan VAE " - "(use --no-taehv for the Wan VAE path)") - parser.add_argument("--taehv_checkpoint", default=DEFAULT_TAEHV_CHECKPOINT, + "(use --no-taehv for the Wan VAE path)") + parser.add_argument("--taehv_checkpoint", + default=DEFAULT_TAEHV_CHECKPOINT, help="Path to the TAEHV taew2_1.pth checkpoint") - parser.add_argument("--distilled_model", default=DEFAULT_DISTILLED_MODEL, + parser.add_argument("--distilled_model", + default=DEFAULT_DISTILLED_MODEL, help="HuggingFace ID (or local path) of a distilled " - "transformer checkpoint to load on top of --model. " - "Pass '' to use the base --model weights instead.") + "transformer checkpoint to load on top of --model. " + "Pass '' to use the base --model weights instead.") parser.add_argument("--num_gpus", type=int, default=1) parser.add_argument("--infer_steps", type=int, default=3) parser.add_argument("--guidance_scale", type=float, default=1.0) - parser.add_argument("--warmups", type=int, default=5, - help="Warmup runs before timing (default: 5).") - parser.add_argument("--benchmark-runs", type=int, default=20, + parser.add_argument("--warmups", type=int, default=5, help="Warmup runs before timing (default: 5).") + parser.add_argument("--benchmark-runs", + type=int, + default=20, help="Timed runs to collect min/max/mean/std over (default: 20).") args = parser.parse_args() @@ -237,11 +230,18 @@ def main() -> None: # runs below measure steady-state latency only. with silence_request_log(): for _ in range(args.warmups): - warm = generator.generate(request={ - "prompt": PROMPT, - "sampling": {"num_inference_steps": 2, "guidance_scale": args.guidance_scale}, - "output": {"save_video": False, "return_frames": args.taehv}, - }) + warm = generator.generate( + request={ + "prompt": PROMPT, + "sampling": { + "num_inference_steps": 2, + "guidance_scale": args.guidance_scale + }, + "output": { + "save_video": False, + "return_frames": args.taehv + }, + }) if args.taehv: taehv.decode(warm.samples) @@ -257,18 +257,19 @@ def main() -> None: frames = None with silence_request_log(): for i in range(args.benchmark_runs): - result = generator.generate(request={ - "prompt": PROMPT, - "sampling": { - "num_inference_steps": args.infer_steps, - "guidance_scale": args.guidance_scale, - }, - "output": { - "save_video": False, - "return_frames": args.taehv, - "output_path": output_path, - }, - }) + result = generator.generate( + request={ + "prompt": PROMPT, + "sampling": { + "num_inference_steps": args.infer_steps, + "guidance_scale": args.guidance_scale, + }, + "output": { + "save_video": False, + "return_frames": args.taehv, + "output_path": output_path, + }, + }) denoise_elapsed = result.generation_time denoise_times.append(denoise_elapsed) diff --git a/examples/inference/optimizations/attention_example.py b/examples/inference/optimizations/attention_example.py index 558a001d2..ba8e830a5 100644 --- a/examples/inference/optimizations/attention_example.py +++ b/examples/inference/optimizations/attention_example.py @@ -3,8 +3,9 @@ import time from fastvideo import VideoGenerator + def main(): - # set the attention backend + # set the attention backend os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" start_time = time.perf_counter() @@ -14,7 +15,7 @@ def main(): dit_cpu_offload=False, vae_cpu_offload=False, text_encoder_cpu_offload=True, - pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" + pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" ) load_time = time.perf_counter() - start_time print(f"Model loading time: {load_time:.2f} seconds") @@ -26,12 +27,13 @@ def main(): "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.", seed=1024, output_path="example_outputs/") - + generation_time = time.perf_counter() - gen_start_time print(f"Video generation time: {generation_time:.2f} seconds") total_time = time.perf_counter() - start_time print(f"Total execution time: {total_time:.2f} seconds") + if __name__ == "__main__": main() diff --git a/examples/inference/optimizations/fp4_attn_wan2_1_1_3b.py b/examples/inference/optimizations/fp4_attn_wan2_1_1_3b.py index 8e7dfa87c..1221e27c1 100644 --- a/examples/inference/optimizations/fp4_attn_wan2_1_1_3b.py +++ b/examples/inference/optimizations/fp4_attn_wan2_1_1_3b.py @@ -24,12 +24,9 @@ OUTPUT_PATH = "video_samples" def main(): parser = argparse.ArgumentParser(description="FP4 FA4 video generation benchmark") - parser.add_argument("--nvfp4_fa4", action="store_true", - help="Enable NVFP4 FP4 quantized QK flash attention") - parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - help="Model path or HuggingFace ID") - parser.add_argument("--compile", action="store_true", - help="Enable torch.compile for DIT") + parser.add_argument("--nvfp4_fa4", action="store_true", help="Enable NVFP4 FP4 quantized QK flash attention") + parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", help="Model path or HuggingFace ID") + parser.add_argument("--compile", action="store_true", help="Enable torch.compile for DIT") parser.add_argument("--num_gpus", type=int, default=1) parser.add_argument("--infer_steps", type=int, default=50) args = parser.parse_args() @@ -51,24 +48,35 @@ def main(): enable_torch_compile=args.compile, ) - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") n_warmup = 2 if args.compile else 1 for i in range(n_warmup): - generator.generate(request={"prompt": prompt, "sampling": {"num_inference_steps": 2}, - "output": {"save_video": False}}) + generator.generate(request={ + "prompt": prompt, + "sampling": { + "num_inference_steps": 2 + }, + "output": { + "save_video": False + } + }) os.makedirs(OUTPUT_PATH, exist_ok=True) start = time.time() - generator.generate(request={ - "prompt": prompt, - "sampling": {"num_inference_steps": args.infer_steps}, - "output": {"save_video": True, "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")}, - }) + generator.generate( + request={ + "prompt": prompt, + "sampling": { + "num_inference_steps": args.infer_steps + }, + "output": { + "save_video": True, + "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4") + }, + }) elapsed = time.time() - start print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.2f}s " f"({args.infer_steps / elapsed:.2f} it/s)") diff --git a/examples/inference/optimizations/fp8_wan2_1_1_3b.py b/examples/inference/optimizations/fp8_wan2_1_1_3b.py index e48d85cc8..735aba6de 100644 --- a/examples/inference/optimizations/fp8_wan2_1_1_3b.py +++ b/examples/inference/optimizations/fp8_wan2_1_1_3b.py @@ -41,8 +41,7 @@ def load_taehv(checkpoint_path, device="cuda", dtype=torch.float16): @torch.no_grad() # type: ignore[misc] def decode_with_taehv(taehv_model, latents): latents = latents.permute(0, 2, 1, 3, 4) - latents = latents.to(device=next(taehv_model.parameters()).device, - dtype=next(taehv_model.parameters()).dtype) + latents = latents.to(device=next(taehv_model.parameters()).device, dtype=next(taehv_model.parameters()).dtype) decoded = taehv_model.decode_video(latents, parallel=False, show_progress_bar=False) frames = [] for frame in decoded[0]: @@ -53,14 +52,16 @@ def decode_with_taehv(taehv_model, latents): def main(): parser = argparse.ArgumentParser(description="FP8 video generation benchmark") - parser.add_argument("--bf16", action="store_true", - help="BF16 baseline (no FP8 quantization)") - parser.add_argument("--granularity", choices=["tensor", "channel"], default="tensor", + parser.add_argument("--bf16", action="store_true", help="BF16 baseline (no FP8 quantization)") + parser.add_argument("--granularity", + choices=["tensor", "channel"], + default="tensor", help="FP8 weight scale granularity: tensor (faster) or channel (more accurate)") - parser.add_argument("--taehv-checkpoint", default=None, metavar="PATH", + parser.add_argument("--taehv-checkpoint", + default=None, + metavar="PATH", help="Path to taew2_1.pth; enables TAEHV tiny autoencoder decoding") - parser.add_argument("--model", default="FastVideo/FastWan-QAD-FP8-1.3B", - help="Model path or HuggingFace ID") + parser.add_argument("--model", default="FastVideo/FastWan-QAD-FP8-1.3B", help="Model path or HuggingFace ID") parser.add_argument("--no-compile", action="store_true", help="Disable torch.compile for the DiT") parser.add_argument("--num_gpus", type=int, default=1) parser.add_argument("--infer_steps", type=int, default=3) @@ -81,9 +82,7 @@ def main(): # transformer_quant needs a QuantizationConfig *instance* — the bare string # is not resolved on the from_pretrained kwarg path. - extra = {} if args.bf16 else { - "transformer_quant": get_quantization_config("FP8")(granularity=args.granularity) - } + extra = {} if args.bf16 else {"transformer_quant": get_quantization_config("FP8")(granularity=args.granularity)} generator = VideoGenerator.from_pretrained( args.model, num_gpus=args.num_gpus, @@ -99,36 +98,56 @@ def main(): **extra, ) - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") n_warmup = 1 if not args.no_compile else 0 for _ in range(n_warmup): - generator.generate(request={"prompt": prompt, "sampling": {"num_inference_steps": 3, "guidance_scale": 1.0}, - "output": {"save_video": False}}) + generator.generate( + request={ + "prompt": prompt, + "sampling": { + "num_inference_steps": 3, + "guidance_scale": 1.0 + }, + "output": { + "save_video": False + } + }) os.makedirs(OUTPUT_PATH, exist_ok=True) start = time.time() if use_taehv: - result = generator.generate(request={ - "prompt": prompt, - "sampling": {"num_inference_steps": args.infer_steps, "guidance_scale": 1.0}, - "output": {"save_video": False}, - }) + result = generator.generate( + request={ + "prompt": prompt, + "sampling": { + "num_inference_steps": args.infer_steps, + "guidance_scale": 1.0 + }, + "output": { + "save_video": False + }, + }) import imageio frames = decode_with_taehv(taehv_model, result.samples) video_path = os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4") imageio.mimsave(video_path, frames, fps=16, format="mp4") print(f"Saved TAEHV-decoded video to: {video_path}") else: - generator.generate(request={ - "prompt": prompt, - "sampling": {"num_inference_steps": args.infer_steps, "guidance_scale": 1.0}, - "output": {"save_video": True, "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")}, - }) + generator.generate( + request={ + "prompt": prompt, + "sampling": { + "num_inference_steps": args.infer_steps, + "guidance_scale": 1.0 + }, + "output": { + "save_video": True, + "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4") + }, + }) elapsed = time.time() - start print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.2f}s " f"({args.infer_steps / elapsed:.2f} it/s)") diff --git a/examples/inference/optimizations/nvfp4_qat_wan2_1_1_3b.py b/examples/inference/optimizations/nvfp4_qat_wan2_1_1_3b.py index 095c8bba4..c4bf75061 100644 --- a/examples/inference/optimizations/nvfp4_qat_wan2_1_1_3b.py +++ b/examples/inference/optimizations/nvfp4_qat_wan2_1_1_3b.py @@ -25,14 +25,14 @@ OUTPUT_PATH = "video_samples" def main(): parser = argparse.ArgumentParser(description="NVFP4 + Attn-QAT video generation") - parser.add_argument("--bf16", action="store_true", - help="BF16 baseline (no NVFP4 linear, default attention)") - parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - help="Model path or HuggingFace ID") - parser.add_argument("--quant-method", default="nvfp4_qat", choices=["nvfp4_qat", "NVFP4"], + parser.add_argument("--bf16", action="store_true", help="BF16 baseline (no NVFP4 linear, default attention)") + parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", help="Model path or HuggingFace ID") + parser.add_argument("--quant-method", + default="nvfp4_qat", + choices=["nvfp4_qat", "NVFP4"], help="Linear quantization config. Wan-2.1 uses nvfp4_qat (matches its " - "to_q/k/v/out + ffn layers); NVFP4 is LTX2-specific and will NOT " - "quantize Wan.") + "to_q/k/v/out + ffn layers); NVFP4 is LTX2-specific and will NOT " + "quantize Wan.") parser.add_argument("--compile", action="store_true", help="Enable torch.compile for the DiT") parser.add_argument("--num_gpus", type=int, default=1) parser.add_argument("--infer_steps", type=int, default=50) @@ -66,24 +66,35 @@ def main(): **extra, ) - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") n_warmup = 2 if args.compile else 1 for _ in range(n_warmup): - generator.generate(request={"prompt": prompt, "sampling": {"num_inference_steps": 2}, - "output": {"save_video": False}}) + generator.generate(request={ + "prompt": prompt, + "sampling": { + "num_inference_steps": 2 + }, + "output": { + "save_video": False + } + }) os.makedirs(OUTPUT_PATH, exist_ok=True) start = time.time() - generator.generate(request={ - "prompt": prompt, - "sampling": {"num_inference_steps": args.infer_steps}, - "output": {"save_video": True, "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")}, - }) + generator.generate( + request={ + "prompt": prompt, + "sampling": { + "num_inference_steps": args.infer_steps + }, + "output": { + "save_video": True, + "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4") + }, + }) elapsed = time.time() - start print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.2f}s " f"({args.infer_steps / elapsed:.2f} it/s)") diff --git a/examples/inference/optimizations/qad_fp4_ab.py b/examples/inference/optimizations/qad_fp4_ab.py index 8dd09d9c4..4d2457609 100644 --- a/examples/inference/optimizations/qad_fp4_ab.py +++ b/examples/inference/optimizations/qad_fp4_ab.py @@ -69,11 +69,9 @@ DEFAULT_DISTILLED = "FastVideo/FastWan-QAD-1.3B" # base Wan pipeline (vae/text_encoder are Wan-identical). DISTILLED_WEIGHTS_FILE = "transformer/diffusion_pytorch_model.safetensors" -PROMPT = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." -) +PROMPT = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") def resolve_distilled_weights(hf_id: str) -> str: @@ -95,7 +93,7 @@ def build_generator(fp4_linear: bool): pipeline_config = PipelineConfig.from_pretrained(model_id) pipeline_config.dit_precision = "bf16" pipeline_config.vae_precision = "bf16" - pipeline_config.text_encoder_precisions = ("bf16",) + pipeline_config.text_encoder_precisions = ("bf16", ) if fp4_linear: # Wan-style config: matches to_q/k/v/out + ffn (the plain NVFP4 config is @@ -159,19 +157,23 @@ def main() -> None: def _generate(): # seed + frame dims live under `sampling` (SamplingConfig); `output` # only takes output_path/save_video/return_frames (OutputConfig). - return generator.generate(request={ - "prompt": prompt, - "sampling": { - "seed": seed, - "num_inference_steps": steps, - "guidance_scale": guidance, - "height": height, - "width": width, - "num_frames": frames, - }, - "output": {"save_video": True, "output_path": arm_dir, - "return_frames": True}, - }) + return generator.generate( + request={ + "prompt": prompt, + "sampling": { + "seed": seed, + "num_inference_steps": steps, + "guidance_scale": guidance, + "height": height, + "width": width, + "num_frames": frames, + }, + "output": { + "save_video": True, + "output_path": arm_dir, + "return_frames": True + }, + }) for _ in range(warmup): _generate() diff --git a/examples/inference/optimizations/spark_benchmark.py b/examples/inference/optimizations/spark_benchmark.py index 0850a6494..e048b5e8f 100644 --- a/examples/inference/optimizations/spark_benchmark.py +++ b/examples/inference/optimizations/spark_benchmark.py @@ -81,17 +81,14 @@ def bench_generation(args) -> None: sampling_param.num_inference_steps = args.steps sampling_param.seed = args.seed - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, " - "its eyes wide with interest. Soft natural light, warm cheerful tones, " - "mid-shot, cinematic.") + prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, " + "its eyes wide with interest. Soft natural light, warm cheerful tones, " + "mid-shot, cinematic.") def _gen(): torch.cuda.synchronize() t0 = time.perf_counter() - video = generator.generate_video(prompt, output_path=args.out, - save_video=True, - sampling_param=sampling_param) + video = generator.generate_video(prompt, output_path=args.out, save_video=True, sampling_param=sampling_param) torch.cuda.synchronize() # generate_video returns a plain dict (legacy result), not an object — # attribute access would silently fall back to wall time / None. @@ -132,28 +129,24 @@ def bench_decode_precision(args) -> None: """Part B: fp32-vs-bf16 VAE decode of one fixed latent (SSIM + speedup).""" try: from diffusers import AutoencoderKLWan - from torchmetrics.functional import ( - multiscale_structural_similarity_index_measure as msssim) + from torchmetrics.functional import (multiscale_structural_similarity_index_measure as msssim) except ImportError as e: # torchmetrics is not a hard FastVideo dep _p(f"skipping decode A/B (missing dependency: {e}); " "`uv pip install torchmetrics` to enable it.") return dev = "cuda" - vae = AutoencoderKLWan.from_pretrained( - args.model, subfolder="vae", torch_dtype=torch.float32).to(dev).eval() + vae = AutoencoderKLWan.from_pretrained(args.model, subfolder="vae", torch_dtype=torch.float32).to(dev).eval() # Wan latent geometry for height x width x frames, patch (4,8,8): # T_lat = (frames - 1) // 4 + 1, H_lat = height // 8, W_lat = width // 8 t_lat = (args.frames - 1) // 4 + 1 - z = torch.randn(1, 16, t_lat, args.height // 8, args.width // 8, - device=dev, dtype=torch.float32) + z = torch.randn(1, 16, t_lat, args.height // 8, args.width // 8, device=dev, dtype=torch.float32) def _decode(autocast: bool): torch.cuda.synchronize() t0 = time.perf_counter() - with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16, - enabled=autocast): + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast): out = vae.decode(z, return_dict=False)[0] torch.cuda.synchronize() return out.float(), time.perf_counter() - t0 @@ -164,8 +157,8 @@ def bench_decode_precision(args) -> None: _decode(False) # warm both paths (excluded from timing) _decode(True) - o32, t32 = _decode(autocast=False) # fp32 - o16, t16 = _decode(autocast=True) # bf16 (== vae_decode_precision="bf16") + o32, t32 = _decode(autocast=False) # fp32 + o16, t16 = _decode(autocast=True) # bf16 (== vae_decode_precision="bf16") ssim = msssim(_frames(o16), _frames(o32), data_range=1.0).item() _p(f"fp32 decode : {t32 * 1000:8.1f} ms") @@ -185,10 +178,8 @@ def main() -> None: ap.add_argument("--width", type=int, default=832) ap.add_argument("--seed", type=int, default=42) ap.add_argument("--out", default="spark_bench_samples") - ap.add_argument("--skip-gen", action="store_true", - help="skip Part A (generation timing)") - ap.add_argument("--skip-decode", action="store_true", - help="skip Part B (decode precision A/B)") + ap.add_argument("--skip-gen", action="store_true", help="skip Part A (generation timing)") + ap.add_argument("--skip-decode", action="store_true", help="skip Part B (decode precision A/B)") args = ap.parse_args() if not torch.cuda.is_available(): diff --git a/examples/inference/optimizations/text_encoder_quant_example.py b/examples/inference/optimizations/text_encoder_quant_example.py index 8b3cca506..b0e22f233 100644 --- a/examples/inference/optimizations/text_encoder_quant_example.py +++ b/examples/inference/optimizations/text_encoder_quant_example.py @@ -18,7 +18,7 @@ def main(text_encoder_path: str): dit_cpu_offload=True, vae_cpu_offload=False, text_encoder_cpu_offload=False, - # AbsMaxFP8 is the quantization method used by ComfyUI; + # AbsMaxFP8 is the quantization method used by ComfyUI; # check fastvideo/layers/quantization/* for more quantization methods override_text_encoder_quant="AbsMaxFP8", # for Wan 2.2, this is the path to "umt5_xxl_fp8_e4m3fn_scaled.safetensors" @@ -29,9 +29,7 @@ def main(text_encoder_path: str): # I2V is triggered just by passing in an image_path argument prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG" - video = generator.generate_video( - prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path - ) + video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path) if __name__ == "__main__": diff --git a/examples/inference/optimizations/torch_compile_example.py b/examples/inference/optimizations/torch_compile_example.py index 2a7cfbf13..bcf47f205 100644 --- a/examples/inference/optimizations/torch_compile_example.py +++ b/examples/inference/optimizations/torch_compile_example.py @@ -23,17 +23,14 @@ import time from fastvideo import VideoGenerator -PROMPT = ( - "A high-definition video of a robotic arm welding a metal structure, " - "bright sparks and smoke, industrial setting." -) +PROMPT = ("A high-definition video of a robotic arm welding a metal structure, " + "bright sparks and smoke, industrial setting.") def main() -> None: parser = argparse.ArgumentParser(description="torch.compile A/B") parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers") - parser.add_argument("--compile", action="store_true", - help="Enable torch.compile for the DiT") + parser.add_argument("--compile", action="store_true", help="Enable torch.compile for the DiT") parser.add_argument("--num_gpus", type=int, default=1) args = parser.parse_args() @@ -54,12 +51,15 @@ def main() -> None: # prompt/seed/shapes both runs so the compiled graph is reused. request: dict = { "prompt": PROMPT, - "sampling": {"seed": 1024}, - "output": {"save_video": save}, + "sampling": { + "seed": 1024 + }, + "output": { + "save_video": save + }, } if save: - request["output"]["output_path"] = ( - f"video_samples/torch_compile_{tag}.mp4") + request["output"]["output_path"] = (f"video_samples/torch_compile_{tag}.mp4") t0 = time.perf_counter() generator.generate(request) return time.perf_counter() - t0 diff --git a/examples/train/scenario/ode_init_self_forcing_wan_causal/decode_kd_cache.py b/examples/train/scenario/ode_init_self_forcing_wan_causal/decode_kd_cache.py index 6e94c3b09..182ab40c8 100644 --- a/examples/train/scenario/ode_init_self_forcing_wan_causal/decode_kd_cache.py +++ b/examples/train/scenario/ode_init_self_forcing_wan_causal/decode_kd_cache.py @@ -28,15 +28,12 @@ def denormalize(latents: torch.Tensor, vae: AutoencoderKLWan) -> torch.Tensor: inverse: latent = normalized * std + mean """ cfg = vae.config - mean = torch.tensor(cfg.latents_mean, dtype=latents.dtype, - device=latents.device).view(1, -1, 1, 1, 1) - std = torch.tensor(cfg.latents_std, dtype=latents.dtype, - device=latents.device).view(1, -1, 1, 1, 1) + mean = torch.tensor(cfg.latents_mean, dtype=latents.dtype, device=latents.device).view(1, -1, 1, 1, 1) + std = torch.tensor(cfg.latents_std, dtype=latents.dtype, device=latents.device).view(1, -1, 1, 1, 1) return latents * std + mean -def decode_real(pt_path: Path, vae: AutoencoderKLWan, - device: torch.device) -> np.ndarray: +def decode_real(pt_path: Path, vae: AutoencoderKLWan, device: torch.device) -> np.ndarray: """Load one .pt cache file and decode its 'real' latent to uint8 [T,H,W,3].""" d = torch.load(pt_path, weights_only=True, map_location="cpu") real = d["real"].float() # [T, C, H, W] (normalized) @@ -69,8 +66,7 @@ def main() -> None: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Loading VAE from {args.model_id} ...") - vae = AutoencoderKLWan.from_pretrained( - args.model_id, subfolder="vae", torch_dtype=torch.float32) + vae = AutoencoderKLWan.from_pretrained(args.model_id, subfolder="vae", torch_dtype=torch.float32) vae.eval().to(device) pts = sorted(samples_dir.glob("*.pt")) diff --git a/fastvideo-kernel/attn_qat_infer/api.py b/fastvideo-kernel/attn_qat_infer/api.py index 088d1e5ed..f506d692d 100644 --- a/fastvideo-kernel/attn_qat_infer/api.py +++ b/fastvideo-kernel/attn_qat_infer/api.py @@ -30,27 +30,21 @@ BLOCK_N = 128 # Block size for N dimension (key/value sequence length) @triton.jit -def group_mean_kernel( - q_ptr, - q_out_ptr, - qm_out_ptr, - B, H, L, D: tl.constexpr, - stride_qb, stride_qh, stride_ql, stride_qd, - stride_qmb, stride_qmh, stride_qml, stride_qmd, - GROUP_SIZE: tl.constexpr -): +def group_mean_kernel(q_ptr, q_out_ptr, qm_out_ptr, B, H, L, D: tl.constexpr, stride_qb, stride_qh, stride_ql, + stride_qd, stride_qmb, stride_qmh, stride_qml, stride_qmd, GROUP_SIZE: tl.constexpr): pid_b = tl.program_id(0) pid_h = tl.program_id(1) pid_group = tl.program_id(2) - + group_start = pid_group * GROUP_SIZE offsets = group_start + tl.arange(0, GROUP_SIZE) - - q_offsets = pid_b * stride_qb + pid_h * stride_qh + offsets[:, None] * stride_ql + tl.arange(0, D)[None, :] * stride_qd + + q_offsets = pid_b * stride_qb + pid_h * stride_qh + offsets[:, None] * stride_ql + tl.arange(0, + D)[None, :] * stride_qd q_group = tl.load(q_ptr + q_offsets) - + qm_group = tl.sum(q_group, axis=0) / GROUP_SIZE - + q_group = q_group - qm_group tl.store(q_out_ptr + q_offsets, q_group) @@ -62,23 +56,37 @@ def triton_group_mean(q: torch.Tensor): B, H, L, D = q.shape GROUP_SIZE = BLOCK_M num_groups = L // GROUP_SIZE - + q_out = torch.empty_like(q) # [B, H, L, D] - qm = torch.empty(B, H, num_groups, D, device=q.device, dtype=q.dtype) - + qm = torch.empty(B, H, num_groups, D, device=q.device, dtype=q.dtype) + grid = (B, H, num_groups) - - group_mean_kernel[grid]( - q, q_out, qm, - B, H, L, D, - q.stride(0), q.stride(1), q.stride(2), q.stride(3), - qm.stride(0), qm.stride(1), qm.stride(2), qm.stride(3), - GROUP_SIZE=GROUP_SIZE - ) + + group_mean_kernel[grid](q, + q_out, + qm, + B, + H, + L, + D, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + qm.stride(0), + qm.stride(1), + qm.stride(2), + qm.stride(3), + GROUP_SIZE=GROUP_SIZE) return q_out, qm -def preprocess_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, per_block_mean: bool = True, enable_smoothing_q: bool = False, enable_smoothing_k: bool = False): +def preprocess_qkv(q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + per_block_mean: bool = True, + enable_smoothing_q: bool = False, + enable_smoothing_k: bool = False): def pad_to_block_size(x): L = x.size(2) @@ -86,9 +94,9 @@ def preprocess_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, per_block_ if pad_len == 0: return x.contiguous() return F.pad(x, (0, 0, 0, pad_len), value=0).contiguous() - + if enable_smoothing_k: - k -= k.mean(dim=-2, keepdim=True) + k -= k.mean(dim=-2, keepdim=True) q, k, v = map(lambda x: pad_to_block_size(x), [q, k, v]) if per_block_mean and enable_smoothing_q: q, qm = triton_group_mean(q) @@ -100,9 +108,10 @@ def preprocess_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, per_block_ else: # used to disable q smoothing B, H, L, D = q.shape delta_s = torch.zeros((B, H, L // BLOCK_M, k.shape[2]), device=q.device, dtype=torch.float32) - + return q, k, v, delta_s + def scale_and_quant_fp4(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: assert x.ndim == 4 B, H, N, D = x.shape @@ -111,6 +120,7 @@ def scale_and_quant_fp4(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: fp4quant_cuda.scaled_fp4_quant(x, packed_fp4, fp8_scale, 1) return packed_fp4, fp8_scale + def scale_and_quant_fp4_permute(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: assert x.ndim == 4 B, H, N, D = x.shape @@ -119,6 +129,7 @@ def scale_and_quant_fp4_permute(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Te fp4quant_cuda.scaled_fp4_quant_permute(x, packed_fp4, fp8_scale, 1) return packed_fp4, fp8_scale + def scale_and_quant_fp4_transpose(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: assert x.ndim == 4 B, H, N, D = x.shape @@ -127,22 +138,31 @@ def scale_and_quant_fp4_transpose(x: torch.Tensor) -> Tuple[torch.Tensor, torch. fp4quant_cuda.scaled_fp4_quant_trans(x, packed_fp4, fp8_scale, 1) return packed_fp4, fp8_scale -def blockscaled_fp4_attn(qlist: Tuple, + +def blockscaled_fp4_attn(qlist: Tuple, klist: Tuple, vlist: Tuple, delta_s: torch.Tensor, KL: int, - is_causal: bool = False, + is_causal: bool = False, per_block_mean: bool = True, is_bf16: bool = True, single_level_p_quant: bool = False, - sm_scale: float | None = None - ): - softmax_scale = sm_scale if sm_scale is not None else (qlist[0].shape[-1] * 2) ** (-0.5) - return fp4attn_cuda.fwd(qlist[0], klist[0], vlist[0], qlist[1], klist[1], vlist[1], delta_s, KL, None, softmax_scale, is_causal, per_block_mean, is_bf16, single_level_p_quant) + sm_scale: float | None = None): + softmax_scale = sm_scale if sm_scale is not None else (qlist[0].shape[-1] * 2)**(-0.5) + return fp4attn_cuda.fwd(qlist[0], klist[0], vlist[0], qlist[1], klist[1], vlist[1], delta_s, KL, None, + softmax_scale, is_causal, per_block_mean, is_bf16, single_level_p_quant) -def sageattn_blackwell(q, k, v, attn_mask = None, is_causal = False, per_block_mean = True, single_level_p_quant = True, sm_scale: float | None = None, **kwargs): +def sageattn_blackwell(q, + k, + v, + attn_mask=None, + is_causal=False, + per_block_mean=True, + single_level_p_quant=True, + sm_scale: float | None = None, + **kwargs): """ SageAttention3 Blackwell kernel for FP4 attention. @@ -166,7 +186,7 @@ def sageattn_blackwell(q, k, v, attn_mask = None, is_causal = False, per_block_m """ if q.size(-1) >= 256: print(f"Unsupported Headdim {q.size(-1)}") - return sdpa(q, k, v, is_causal = is_causal) + return sdpa(q, k, v, is_causal=is_causal) QL = q.size(2) KL = k.size(2) is_bf16 = q.dtype == torch.bfloat16 @@ -174,16 +194,6 @@ def sageattn_blackwell(q, k, v, attn_mask = None, is_causal = False, per_block_m qlist_from_cuda = scale_and_quant_fp4(q) klist_from_cuda = scale_and_quant_fp4_permute(k) vlist_from_cuda = scale_and_quant_fp4_transpose(v) - o_fp4 = blockscaled_fp4_attn( - qlist_from_cuda, - klist_from_cuda, - vlist_from_cuda, - delta_s, - KL, - is_causal, - per_block_mean, - is_bf16, - single_level_p_quant, - sm_scale - )[0][:, :, :QL, :].contiguous() + o_fp4 = blockscaled_fp4_attn(qlist_from_cuda, klist_from_cuda, vlist_from_cuda, delta_s, KL, is_causal, + per_block_mean, is_bf16, single_level_p_quant, sm_scale)[0][:, :, :QL, :].contiguous() return o_fp4 diff --git a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_k.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_k.py index 94c10681c..5fc97f5ee 100644 --- a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_k.py +++ b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_k.py @@ -17,18 +17,21 @@ import torch import fp4quant from triton.tools.mxfp import MXFP4Tensor -from bench_utils import bench_kineto +from bench_utils import bench_kineto + b = 1 h = 32 n = 16384 d = 128 + def test(): q = torch.randn((b, h, n, d), device="cuda", dtype=torch.float16) o = torch.empty((b, h, n, d // 2), device="cuda", dtype=torch.uint8) o_s = torch.empty((b, h, n, d // 16), device="cuda", dtype=torch.float8_e4m3fn) fp4quant.scaled_fp4_quant_permute(q, o, o_s, 1) + test() t = bench_kineto(test, "scaled_fp4_quant_kernel", suppress_kineto_output=True) @@ -38,23 +41,27 @@ throughput = IO / t * 1e-9 print(f"Throughput: {throughput:.2f} GB/s") + def scale_and_fp4_tensor(x: torch.Tensor, packed_dim: int = 3, all_ones: bool = False, permuted: bool = False): assert x.is_contiguous() and x.ndim == 4 and x.shape[-1] % 16 == 0 B, H, M, N = x.shape x = x.view(B, H, M, N // 16, 16) scales = (x.abs().amax(dim=-1, keepdim=True) / 6).to(torch.float32) if all_ones: - scales = torch.ones_like(scales) + scales = torch.ones_like(scales) x_scaled = x / scales packed_fp4 = MXFP4Tensor(x_scaled.flatten(start_dim=-2)).to_packed_tensor(dim=packed_dim) - dequant_x = (MXFP4Tensor(x_scaled).to(torch.float32) * scales.to(torch.float8_e4m3fn).to(torch.float32)).flatten(start_dim=-2) + dequant_x = (MXFP4Tensor(x_scaled).to(torch.float32) * + scales.to(torch.float8_e4m3fn).to(torch.float32)).flatten(start_dim=-2) fp8_scale = scales.flatten(start_dim=-2).to(torch.float8_e4m3fn) permuted_fp8_scale = None if permuted: scales = scales.view(B, H // 64, 4, 16, M, N // 16).permute(0, 1, 3, 2, 4, 5).reshape(B, H, M, N // 16) - permuted_fp8_scale = scales.view(B, H // 64, 64, M, N // 64, 4).permute(0, 1, 4, 3, 2, 5).reshape(B, H, M, N // 16).to(torch.float8_e4m3fn) + permuted_fp8_scale = scales.view(B, H // 64, 64, M, N // 64, + 4).permute(0, 1, 4, 3, 2, 5).reshape(B, H, M, N // 16).to(torch.float8_e4m3fn) return fp8_scale, packed_fp4, dequant_x, permuted_fp8_scale + b = 2 h = 4 n = 251 @@ -67,7 +74,9 @@ o_s = torch.empty((b, h, n, d // 16), dtype=torch.float8_e4m3fn, device='cuda') fp4quant.scaled_fp4_quant(q, o, o_s, 1) -k_permute = [0, 1, 8, 9, 16, 17, 24, 25, 2, 3, 10, 11, 18, 19, 26, 27, 4, 5, 12, 13, 20, 21, 28, 29, 6, 7, 14, 15, 22, 23, 30, 31] +k_permute = [ + 0, 1, 8, 9, 16, 17, 24, 25, 2, 3, 10, 11, 18, 19, 26, 27, 4, 5, 12, 13, 20, 21, 28, 29, 6, 7, 14, 15, 22, 23, 30, 31 +] o_permuted = torch.empty((b, h, n_padded, d // 2), dtype=torch.uint8, device='cuda') o_s_permuted = torch.empty((b, h, n_padded, d // 16), dtype=torch.float8_e4m3fn, device='cuda') fp4quant.scaled_fp4_quant_permute(q, o_permuted, o_s_permuted, 1) @@ -75,16 +84,19 @@ fp4quant.scaled_fp4_quant_permute(q, o_permuted, o_s_permuted, 1) # padding if n % 128 != 0: o_permuted_gt = torch.cat([o, torch.zeros((b, h, n_padded - n, d // 2), dtype=torch.uint8, device='cuda')], dim=2) - o_s_permuted_gt = torch.cat([o_s, torch.zeros((b, h, n_padded - n, d // 16), dtype=torch.float8_e4m3fn, device='cuda')], dim=2) + o_s_permuted_gt = torch.cat( + [o_s, torch.zeros((b, h, n_padded - n, d // 16), dtype=torch.float8_e4m3fn, device='cuda')], dim=2) else: o_permuted_gt = o o_s_permuted_gt = o_s # use scale_and_fp4_tensor + torch permutation to get the ground truth -o_permuted_gt = o_permuted_gt.reshape(b, h, n_padded // 32, 32, d // 2)[:, :, :, k_permute, :].reshape(b, h, n_padded, d // 2) -o_s_permuted_gt = o_s_permuted_gt.reshape(b, h, n_padded // 32, 32, d // 16)[:, :, :, k_permute, :].reshape(b, h, n_padded, d // 16) +o_permuted_gt = o_permuted_gt.reshape(b, h, n_padded // 32, 32, d // 2)[:, :, :, + k_permute, :].reshape(b, h, n_padded, d // 2) +o_s_permuted_gt = o_s_permuted_gt.reshape(b, h, n_padded // 32, 32, + d // 16)[:, :, :, k_permute, :].reshape(b, h, n_padded, d // 16) -assert((o_permuted - o_permuted_gt).abs().max() == 0) -assert((o_s_permuted.float() - o_s_permuted_gt.float()).abs().max() == 0) +assert ((o_permuted - o_permuted_gt).abs().max() == 0) +assert ((o_s_permuted.float() - o_s_permuted_gt.float()).abs().max() == 0) -print("All tests passed!") \ No newline at end of file +print("All tests passed!") diff --git a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_q.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_q.py index ffbab0a7b..1d0576d59 100644 --- a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_q.py +++ b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_q.py @@ -17,18 +17,21 @@ import torch import fp4quant from triton.tools.mxfp import MXFP4Tensor -from bench_utils import bench_kineto +from bench_utils import bench_kineto + b = 1 h = 32 n = 16384 d = 128 + def test(): q = torch.randn((b, h, n, d), device="cuda", dtype=torch.float16) o = torch.empty((b, h, n, d // 2), device="cuda", dtype=torch.uint8) o_s = torch.empty((b, h, n, d // 16), device="cuda", dtype=torch.float8_e4m3fn) fp4quant.scaled_fp4_quant(q, o, o_s, 1) + test() t = bench_kineto(test, "scaled_fp4_quant_kernel", suppress_kineto_output=True) @@ -38,23 +41,27 @@ throughput = IO / t * 1e-9 print(f"Throughput: {throughput:.2f} GB/s") + def scale_and_fp4_tensor(x: torch.Tensor, packed_dim: int = 3, all_ones: bool = False, permuted: bool = False): assert x.is_contiguous() and x.ndim == 4 and x.shape[-1] % 16 == 0 B, H, M, N = x.shape x = x.view(B, H, M, N // 16, 16) scales = (x.abs().amax(dim=-1, keepdim=True) / 6).to(torch.float32) if all_ones: - scales = torch.ones_like(scales) + scales = torch.ones_like(scales) x_scaled = x / scales packed_fp4 = MXFP4Tensor(x_scaled.flatten(start_dim=-2)).to_packed_tensor(dim=packed_dim) - dequant_x = (MXFP4Tensor(x_scaled).to(torch.float32) * scales.to(torch.float8_e4m3fn).to(torch.float32)).flatten(start_dim=-2) + dequant_x = (MXFP4Tensor(x_scaled).to(torch.float32) * + scales.to(torch.float8_e4m3fn).to(torch.float32)).flatten(start_dim=-2) fp8_scale = scales.flatten(start_dim=-2).to(torch.float8_e4m3fn) permuted_fp8_scale = None if permuted: scales = scales.view(B, H // 64, 4, 16, M, N // 16).permute(0, 1, 3, 2, 4, 5).reshape(B, H, M, N // 16) - permuted_fp8_scale = scales.view(B, H // 64, 64, M, N // 64, 4).permute(0, 1, 4, 3, 2, 5).reshape(B, H, M, N // 16).to(torch.float8_e4m3fn) + permuted_fp8_scale = scales.view(B, H // 64, 64, M, N // 64, + 4).permute(0, 1, 4, 3, 2, 5).reshape(B, H, M, N // 16).to(torch.float8_e4m3fn) return fp8_scale, packed_fp4, dequant_x, permuted_fp8_scale + b = 2 h = 4 n = 251 @@ -68,19 +75,14 @@ fp4quant.scaled_fp4_quant(q, o, o_s, 1) fp8_scale, packed_fp4, dequant_x, permuted_fp8_scale = scale_and_fp4_tensor(q, packed_dim=3) -assert((fp8_scale.float() - o_s.float()).abs().max() == 0) +assert ((fp8_scale.float() - o_s.float()).abs().max() == 0) -o_binary = [ - (int(bin_str[:4], 2), int(bin_str[4:], 2)) - for bin_str in [format(x.item(), '08b') for x in o.view(-1)] -] -o_binary_gt = [ - (int(bin_str[:4], 2), int(bin_str[4:], 2)) - for bin_str in [format(x.item(), '08b') for x in packed_fp4.view(-1)] -] +o_binary = [(int(bin_str[:4], 2), int(bin_str[4:], 2)) for bin_str in [format(x.item(), '08b') for x in o.view(-1)]] +o_binary_gt = [(int(bin_str[:4], 2), int(bin_str[4:], 2)) + for bin_str in [format(x.item(), '08b') for x in packed_fp4.view(-1)]] for i in range(len(o_binary)): # check contiguous 4 bits. Difference should be at most one - assert(abs(o_binary[i][0] - o_binary_gt[i][0]) <= 1) - assert(abs(o_binary[i][1] - o_binary_gt[i][1]) <= 1) + assert (abs(o_binary[i][0] - o_binary_gt[i][0]) <= 1) + assert (abs(o_binary[i][1] - o_binary_gt[i][1]) <= 1) -print("All tests passed!") \ No newline at end of file +print("All tests passed!") diff --git a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py index 24b87b700..0592aa888 100644 --- a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py +++ b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py @@ -17,18 +17,21 @@ import torch import fp4quant from triton.tools.mxfp import MXFP4Tensor -from bench_utils import bench_kineto +from bench_utils import bench_kineto + b = 1 h = 32 n = 16384 d = 128 + def test(): q = torch.randn((b, h, n, d), device="cuda", dtype=torch.float16) o = torch.empty((b, h, d, n // 2), device="cuda", dtype=torch.uint8) o_s = torch.empty((b, h, d, n // 16), device="cuda", dtype=torch.float8_e4m3fn) fp4quant.scaled_fp4_quant_trans(q, o, o_s, 1) + test() t = bench_kineto(test, "scaled_fp4_quant_trans_kernel", suppress_kineto_output=True) @@ -38,23 +41,27 @@ throughput = IO / t * 1e-9 print(f"Throughput: {throughput:.2f} GB/s") + def scale_and_fp4_tensor(x: torch.Tensor, packed_dim: int = 3, all_ones: bool = False, permuted: bool = False): assert x.is_contiguous() and x.ndim == 4 and x.shape[-1] % 16 == 0 B, H, M, N = x.shape x = x.view(B, H, M, N // 16, 16) scales = (x.abs().amax(dim=-1, keepdim=True) / 6).to(torch.float32) if all_ones: - scales = torch.ones_like(scales) + scales = torch.ones_like(scales) x_scaled = x / scales packed_fp4 = MXFP4Tensor(x_scaled.flatten(start_dim=-2)).to_packed_tensor(dim=packed_dim) - dequant_x = (MXFP4Tensor(x_scaled).to(torch.float32) * scales.to(torch.float8_e4m3fn).to(torch.float32)).flatten(start_dim=-2) + dequant_x = (MXFP4Tensor(x_scaled).to(torch.float32) * + scales.to(torch.float8_e4m3fn).to(torch.float32)).flatten(start_dim=-2) fp8_scale = scales.flatten(start_dim=-2).to(torch.float8_e4m3fn) permuted_fp8_scale = None if permuted: scales = scales.view(B, H // 64, 4, 16, M, N // 16).permute(0, 1, 3, 2, 4, 5).reshape(B, H, M, N // 16) - permuted_fp8_scale = scales.view(B, H // 64, 64, M, N // 64, 4).permute(0, 1, 4, 3, 2, 5).reshape(B, H, M, N // 16).to(torch.float8_e4m3fn) + permuted_fp8_scale = scales.view(B, H // 64, 64, M, N // 64, + 4).permute(0, 1, 4, 3, 2, 5).reshape(B, H, M, N // 16).to(torch.float8_e4m3fn) return fp8_scale, packed_fp4, dequant_x, permuted_fp8_scale + b = 2 h = 4 n = 491 @@ -72,7 +79,7 @@ if n % 128 != 0: else: q_padded = q -# use torch transpose + scaled_fp4_quant to get the ground truth +# use torch transpose + scaled_fp4_quant to get the ground truth q_padded = q_padded.transpose(2, 3).reshape(b, h, n_padded, d).contiguous() o_gt = torch.empty((b, h, n_padded, d // 2), dtype=torch.uint8, device='cuda') o_s_gt = torch.empty((b, h, n_padded, d // 16), dtype=torch.float8_e4m3fn, device='cuda') @@ -80,7 +87,7 @@ fp4quant.scaled_fp4_quant(q_padded, o_gt, o_s_gt, 1) o_gt = o_gt.reshape(b, h, d, n_padded // 2).contiguous() o_s_gt = o_s_gt.reshape(b, h, d, n_padded // 16).contiguous() -assert((o_s_gt.float() - o_s.float()).abs().max() == 0) -assert((o_gt - o).abs().max() == 0) +assert ((o_s_gt.float() - o_s.float()).abs().max() == 0) +assert ((o_gt - o).abs().max() == 0) -print("All tests passed!") \ No newline at end of file +print("All tests passed!") diff --git a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_utils.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_utils.py index a7046de42..e757efcad 100644 --- a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_utils.py +++ b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_utils.py @@ -19,8 +19,7 @@ import torch import torch.distributed as dist -def bench(fn, num_warmups: int = 5, num_tests: int = 10, - high_precision: bool = False): +def bench(fn, num_warmups: int = 5, num_tests: int = 10, high_precision: bool = False): # Flush L2 cache with 256 MB data torch.cuda.synchronize() cache = torch.empty(int(256e6 // 4), dtype=torch.int, device='cuda') @@ -49,6 +48,7 @@ def bench(fn, num_warmups: int = 5, num_tests: int = 10, class empty_suppress: + def __enter__(self): return self @@ -57,6 +57,7 @@ class empty_suppress: class suppress_stdout_stderr: + def __enter__(self): self.outnull_file = open(os.devnull, 'w') self.errnull_file = open(os.devnull, 'w') @@ -91,8 +92,13 @@ class suppress_stdout_stderr: self.errnull_file.close() -def bench_kineto(fn, kernel_names, num_tests: int = 30, suppress_kineto_output: bool = False, - trace_path: str = None, barrier_comm_profiling: bool = False, flush_l2: bool = False): +def bench_kineto(fn, + kernel_names, + num_tests: int = 30, + suppress_kineto_output: bool = False, + trace_path: str = None, + barrier_comm_profiling: bool = False, + flush_l2: bool = False): # Conflict with Nsight Systems using_nsys = os.environ.get('DG_NSYS_PROFILING', False) @@ -103,7 +109,8 @@ def bench_kineto(fn, kernel_names, num_tests: int = 30, suppress_kineto_output: suppress = suppress_stdout_stderr if suppress_kineto_output and not using_nsys else empty_suppress with suppress(): schedule = torch.profiler.schedule(wait=0, warmup=1, active=1, repeat=1) if not using_nsys else None - profiler = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA], schedule=schedule) if not using_nsys else empty_suppress() + profiler = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA], + schedule=schedule) if not using_nsys else empty_suppress() with profiler: for i in range(2): # NOTES: use a large kernel and a barrier to eliminate the unbalanced CPU launch overhead @@ -166,4 +173,4 @@ def count_bytes(tensors): total += count_bytes(t) else: total += t.numel() * t.element_size() - return total \ No newline at end of file + return total diff --git a/fastvideo-kernel/benchmarks/bench_fused_compress_topk.py b/fastvideo-kernel/benchmarks/bench_fused_compress_topk.py index 79d598a38..d0664caee 100644 --- a/fastvideo-kernel/benchmarks/bench_fused_compress_topk.py +++ b/fastvideo-kernel/benchmarks/bench_fused_compress_topk.py @@ -46,11 +46,15 @@ def parse_arguments() -> argparse.Namespace: p.add_argument("--batch_size", type=int, default=1) p.add_argument("--num_heads", type=int, default=12) p.add_argument("--head_dim", type=int, default=128, choices=[64, 128]) - p.add_argument("--seq_lens", type=int, nargs="+", default=[49152], + p.add_argument("--seq_lens", + type=int, + nargs="+", + default=[49152], help="Sequence lengths to benchmark (must be divisible by block_elements)") - p.add_argument("--block_elements", type=int, default=64, choices=[64, 256], - help="Tokens per block (64 or 256)") - p.add_argument("--topk", type=int, default=None, + p.add_argument("--block_elements", type=int, default=64, choices=[64, 256], help="Tokens per block (64 or 256)") + p.add_argument("--topk", + type=int, + default=None, help="KV blocks to select per Q block (default: ~10%% of kv_blocks)") p.add_argument("--warmup", type=int, default=10) p.add_argument("--rep", type=int, default=50) @@ -63,6 +67,7 @@ def parse_arguments() -> argparse.Namespace: # Old PyTorch baselines (extracted from ops.py before commit d9cd0b9) # --------------------------------------------------------------------------- + def pytorch_block_mean( x: torch.Tensor, variable_block_sizes: torch.Tensor, @@ -91,6 +96,7 @@ def pytorch_topk_mask( # Accuracy helpers # --------------------------------------------------------------------------- + def accuracy_compress(ref: torch.Tensor, test: torch.Tensor) -> dict: ref_f = ref.float() test_f = test.float() @@ -118,13 +124,20 @@ def accuracy_topk(ref_mask: torch.Tensor, test_mask: torch.Tensor) -> dict: # Benchmark runner # --------------------------------------------------------------------------- + def bench_compress_fwd( - B: int, H: int, seq_len: int, D: int, block_elements: int, - dtype: torch.dtype, warmup: int, rep: int, + B: int, + H: int, + seq_len: int, + D: int, + block_elements: int, + dtype: torch.dtype, + warmup: int, + rep: int, ) -> None: num_blocks = seq_len // block_elements x = torch.randn(B, H, seq_len, D, dtype=dtype, device="cuda") - vbs = torch.full((num_blocks,), block_elements, dtype=torch.int32, device="cuda") + vbs = torch.full((num_blocks, ), block_elements, dtype=torch.int32, device="cuda") if num_blocks > 4: vbs[1] = block_elements - 2 vbs[-2] = block_elements - 5 @@ -144,11 +157,17 @@ def bench_compress_fwd( def bench_compress_bwd( - B: int, H: int, seq_len: int, D: int, block_elements: int, - dtype: torch.dtype, warmup: int, rep: int, + B: int, + H: int, + seq_len: int, + D: int, + block_elements: int, + dtype: torch.dtype, + warmup: int, + rep: int, ) -> None: num_blocks = seq_len // block_elements - vbs = torch.full((num_blocks,), block_elements, dtype=torch.int32, device="cuda") + vbs = torch.full((num_blocks, ), block_elements, dtype=torch.int32, device="cuda") if num_blocks > 4: vbs[1] = block_elements - 2 vbs[-2] = block_elements - 5 @@ -176,7 +195,8 @@ def bench_compress_bwd( torch.autograd.grad(loss_o, x_o, retain_graph=True) old_ms = do_bench( lambda: torch.autograd.grad(loss_o, x_o, retain_graph=True), - warmup=0, rep=rep, + warmup=0, + rep=rep, ) x_n = x_old.detach().clone().requires_grad_(True) @@ -186,7 +206,8 @@ def bench_compress_bwd( torch.autograd.grad(loss_n, x_n, retain_graph=True) new_ms = do_bench( lambda: torch.autograd.grad(loss_n, x_n, retain_graph=True), - warmup=0, rep=rep, + warmup=0, + rep=rep, ) speedup = old_ms / new_ms if new_ms > 0 else float("inf") @@ -195,8 +216,13 @@ def bench_compress_bwd( def bench_topk( - B: int, H: int, num_blocks: int, topk: int, - dtype: torch.dtype, warmup: int, rep: int, + B: int, + H: int, + num_blocks: int, + topk: int, + dtype: torch.dtype, + warmup: int, + rep: int, ) -> None: scores = torch.randn(B, H, num_blocks, num_blocks, dtype=dtype, device="cuda") @@ -215,8 +241,12 @@ def bench_topk( def bench_topk_neginf( - B: int, H: int, num_blocks: int, topk: int, - dtype: torch.dtype, inf_ratio: float = 0.3, + B: int, + H: int, + num_blocks: int, + topk: int, + dtype: torch.dtype, + inf_ratio: float = 0.3, ) -> None: """Test topk correctness when a fraction of scores are -inf (masked positions).""" scores = torch.randn(B, H, num_blocks, num_blocks, dtype=dtype, device="cuda") @@ -236,8 +266,13 @@ _MAX_KV_BLOCK_SIZE = 4096 def bench_topk_large_kv( - B: int, H: int, kv_blocks: int, topk: int, - dtype: torch.dtype, warmup: int, rep: int, + B: int, + H: int, + kv_blocks: int, + topk: int, + dtype: torch.dtype, + warmup: int, + rep: int, ) -> None: """Test topk with large kv_blocks that may exceed Triton block size limit.""" import triton diff --git a/fastvideo-kernel/benchmarks/bench_vsa.py b/fastvideo-kernel/benchmarks/bench_vsa.py index 28a9f6aef..5236edc93 100644 --- a/fastvideo-kernel/benchmarks/bench_vsa.py +++ b/fastvideo-kernel/benchmarks/bench_vsa.py @@ -23,7 +23,6 @@ try: except Exception as e: # pragma: no cover raise ImportError("This benchmark requires triton (for triton.testing.do_bench).") from e - BLOCK_M = 64 BLOCK_N = 64 @@ -43,16 +42,23 @@ def parse_arguments() -> argparse.Namespace: p.add_argument("--head_dim", type=int, default=128, choices=[64, 128]) p.add_argument("--topk", type=int, default=None, help="KV blocks per Q block (default: ~90%% sparsity)") p.add_argument("--q_seq_lens", type=int, nargs="+", default=[49152], help="Q sequence lengths (must be /64)") - p.add_argument("--kv_seq_lens", type=int, nargs="+", default=None, help="KV sequence lengths (defaults to q_seq_len)") + p.add_argument("--kv_seq_lens", + type=int, + nargs="+", + default=None, + help="KV sequence lengths (defaults to q_seq_len)") p.add_argument("--warmup", type=int, default=5) p.add_argument("--rep", type=int, default=20) p.add_argument("--seed", type=int, default=42) p.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16"]) - p.add_argument("--force_triton", action="store_true", help="Force wrapper to use Triton path (if supported by shapes).") + p.add_argument("--force_triton", + action="store_true", + help="Force wrapper to use Triton path (if supported by shapes).") return p.parse_args() -def create_qkv(batch: int, heads: int, q_len: int, kv_len: int, d: int, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +def create_qkv(batch: int, heads: int, q_len: int, kv_len: int, d: int, + dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: q = torch.randn(batch, heads, q_len, d, dtype=dtype, device="cuda") k = torch.randn(batch, heads, kv_len, d, dtype=dtype, device="cuda") v = torch.randn(batch, heads, kv_len, d, dtype=dtype, device="cuda") @@ -117,13 +123,14 @@ def main() -> None: topk = min(topk, num_kv_blocks) print("\n" + "=" * 80) - print(f"q_len={q_len}, kv_len={kv_len}, num_q_blocks={num_q_blocks}, num_kv_blocks={num_kv_blocks}, topk={topk}") + print( + f"q_len={q_len}, kv_len={kv_len}, num_q_blocks={num_q_blocks}, num_kv_blocks={num_kv_blocks}, topk={topk}") q, k, v = create_qkv(bs, h, q_len, kv_len, d, dtype) block_map = make_block_map(bs, h, num_q_blocks, num_kv_blocks, topk) # Variable block sizes: default full blocks (64 tokens per KV block) - variable_block_sizes = torch.full((num_kv_blocks,), BLOCK_N, dtype=torch.int32, device="cuda") + variable_block_sizes = torch.full((num_kv_blocks, ), BLOCK_N, dtype=torch.int32, device="cuda") def _fwd(): return block_sparse_attn(q, k, v, block_map, variable_block_sizes) @@ -162,5 +169,3 @@ if __name__ == "__main__": if not torch.cuda.is_available(): raise RuntimeError("CUDA is required for this benchmark.") main() - - diff --git a/fastvideo-kernel/benchmarks/benchmark_attn_qat_train.py b/fastvideo-kernel/benchmarks/benchmark_attn_qat_train.py index 6036fd195..9d7ff4c79 100644 --- a/fastvideo-kernel/benchmarks/benchmark_attn_qat_train.py +++ b/fastvideo-kernel/benchmarks/benchmark_attn_qat_train.py @@ -16,7 +16,6 @@ import torch from fastvideo_kernel.triton_kernels.attn_qat_train import attention - RTX_5090_DENSE_BF16_TFLOPS = 209.5 @@ -68,11 +67,9 @@ def _format_result( median_ms, min_ms, max_ms = timing_ms algorithmic_tflops = algorithmic_flops / (median_ms * 1e9) executed_tflops = executed_matmul_flops / (median_ms * 1e9) - return ( - f"{label}: {median_ms:.3f} ms (min={min_ms:.3f}, max={max_ms:.3f}), " - f"algorithmic={algorithmic_tflops:.2f} TFLOPS/{100 * algorithmic_tflops / peak_tflops:.2f}% MFU, " - f"executed_matmul={executed_tflops:.2f} TFLOPS/{100 * executed_tflops / peak_tflops:.2f}% MFU" - ) + return (f"{label}: {median_ms:.3f} ms (min={min_ms:.3f}, max={max_ms:.3f}), " + f"algorithmic={algorithmic_tflops:.2f} TFLOPS/{100 * algorithmic_tflops / peak_tflops:.2f}% MFU, " + f"executed_matmul={executed_tflops:.2f} TFLOPS/{100 * executed_tflops / peak_tflops:.2f}% MFU") def main() -> None: diff --git a/fastvideo-kernel/python/fastvideo_kernel/__init__.py b/fastvideo-kernel/python/fastvideo_kernel/__init__.py index d8ff068c3..bb398c3a0 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/__init__.py +++ b/fastvideo-kernel/python/fastvideo_kernel/__init__.py @@ -26,8 +26,7 @@ from fastvideo_kernel.turbodiffusion_ops import ( ) from fastvideo_kernel.block_sparse_attn_varlen import ( - block_sparse_attn_varlen, -) + block_sparse_attn_varlen, ) from fastvideo_kernel.vsa_utils import ( VSA_TILE_SIZE, diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py index caad8b672..40617e993 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -7,7 +7,6 @@ from typing import Tuple import torch - # --------------------------------------------------------------------------- # Backend selection helpers # --------------------------------------------------------------------------- @@ -37,10 +36,8 @@ def _force_triton() -> bool: Honors both the new ``FASTVIDEO_VSA_TRITON`` and the legacy ``FASTVIDEO_KERNEL_VSA_FORCE_TRITON`` for backward compatibility. """ - return ( - os.environ.get("FASTVIDEO_VSA_TRITON", "0") == "1" - or os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1" - ) + return (os.environ.get("FASTVIDEO_VSA_TRITON", "0") == "1" + or os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1") def _force_tk() -> bool: @@ -62,25 +59,19 @@ def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: if block_map.dim() == 3: block_map = block_map.unsqueeze(0) if block_map.dim() != 4: - raise ValueError( - f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), " - f"got shape={tuple(block_map.shape)}" - ) + raise ValueError(f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), " + f"got shape={tuple(block_map.shape)}") if block_map.dtype != torch.bool: block_map = block_map.to(torch.bool) if not block_map.is_cuda: - raise RuntimeError( - "block_map must be a CUDA tensor (Triton map_to_index required)." - ) + raise RuntimeError("block_map must be a CUDA tensor (Triton map_to_index required).") try: from fastvideo_kernel.triton_kernels.index import map_to_index as triton_map_to_index except Exception as e: # pragma: no cover - environment issue - raise ImportError( - "Triton map_to_index is required but not available. " - "Ensure Triton is installed and " - "fastvideo_kernel.triton_kernels.index is importable." - ) from e + raise ImportError("Triton map_to_index is required but not available. " + "Ensure Triton is installed and " + "fastvideo_kernel.triton_kernels.index is importable.") from e return triton_map_to_index(block_map) @@ -123,8 +114,7 @@ def block_sparse_attn_triton( variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import ( - triton_block_sparse_attn_forward, - ) + triton_block_sparse_attn_forward, ) o, M = triton_block_sparse_attn_forward( q.contiguous(), @@ -172,13 +162,10 @@ def block_sparse_attn_backward_triton( variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import ( - triton_block_sparse_attn_backward, - ) + triton_block_sparse_attn_backward, ) num_kv_blocks = int(variable_block_sizes.numel()) - k2q_idx, k2q_num = _invert_indices_for_backward( - q2k_idx, q2k_num, num_kv_blocks - ) + k2q_idx, k2q_num = _invert_indices_for_backward(q2k_idx, q2k_num, num_kv_blocks) # q/k/v are saved from the user-facing inputs and may be non-contiguous; # o/M are kernel outputs so are already contiguous. dq, dk, dv = triton_block_sparse_attn_backward( @@ -223,16 +210,11 @@ def _setup_context_triton(ctx, inputs, output): def _backward_triton(ctx, grad_o, grad_M): q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors - dq, dk, dv = block_sparse_attn_backward_triton( - grad_o, q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes - ) + dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes) return dq, dk, dv, None, None, None -block_sparse_attn_triton.register_autograd( - _backward_triton, setup_context=_setup_context_triton -) - +block_sparse_attn_triton.register_autograd(_backward_triton, setup_context=_setup_context_triton) # --------------------------------------------------------------------------- # SM90 backend custom ops (index-native) @@ -306,9 +288,7 @@ def block_sparse_attn_backward_sm90( raise ImportError("fastvideo_kernel_ops.block_sparse_bwd is not available") num_kv_blocks = int(variable_block_sizes.numel()) - k2q_idx, k2q_num = _invert_indices_for_backward( - q2k_idx, q2k_num, num_kv_blocks - ) + k2q_idx, k2q_num = _invert_indices_for_backward(q2k_idx, q2k_num, num_kv_blocks) # q/k/v are saved from user-facing inputs; o/lse are kernel outputs. dq, dk, dv = block_sparse_bwd( @@ -353,16 +333,11 @@ def _setup_context_sm90(ctx, inputs, output): def _backward_sm90(ctx, grad_o, grad_lse): q, k, v, o, lse, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors - dq, dk, dv = block_sparse_attn_backward_sm90( - grad_o, q, k, v, o, lse, q2k_idx, q2k_num, variable_block_sizes - ) + dq, dk, dv = block_sparse_attn_backward_sm90(grad_o, q, k, v, o, lse, q2k_idx, q2k_num, variable_block_sizes) return dq, dk, dv, None, None, None -block_sparse_attn_sm90.register_autograd( - _backward_sm90, setup_context=_setup_context_sm90 -) - +block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_context_sm90) # --------------------------------------------------------------------------- # Public API @@ -385,11 +360,7 @@ def block_sparse_attn_from_indices( variable_block_sizes = _as_int32_contig(variable_block_sizes, "variable_block_sizes") block_sparse_fwd, block_sparse_bwd = _get_sm90_ops() - sm90_available = ( - _is_sm90() - and block_sparse_fwd is not None - and block_sparse_bwd is not None - ) + sm90_available = (_is_sm90() and block_sparse_fwd is not None and block_sparse_bwd is not None) # Backend resolution: # - FASTVIDEO_VSA_TRITON forces Triton everywhere. @@ -419,6 +390,4 @@ def block_sparse_attn( ) -> Tuple[torch.Tensor, torch.Tensor]: """Bool-mask compat wrapper; prefer block_sparse_attn_from_indices.""" q2k_idx, q2k_num = _map_to_index(block_map) - return block_sparse_attn_from_indices( - q, k, v, q2k_idx, q2k_num, variable_block_sizes - ) + return block_sparse_attn_from_indices(q, k, v, q2k_idx, q2k_num, variable_block_sizes) diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py index dc95a5fe6..34c56d6aa 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py @@ -30,7 +30,6 @@ from .block_sparse_attn import block_sparse_attn_triton, _force_triton # FA4 CuTe build (``flash_attn.cute``) and make it a hard dependency of the # default Triton path. - _KV_BLOCK_PHYS = 128 # FA4 CuTe BSA forward uses 128-token KV blocks. _KV_BLOCK_TRITON = 64 # Existing Triton path uses 64-token KV blocks. @@ -66,7 +65,7 @@ def _expand_mask_and_sizes_256_to_128( child0 = torch.clamp(sizes_i32, min=0, max=_KV_BLOCK_PHYS) child1 = torch.clamp(sizes_i32 - _KV_BLOCK_PHYS, min=0, max=_KV_BLOCK_PHYS) expanded_sizes = torch.empty( - (sizes_i32.numel() * 2,), + (sizes_i32.numel() * 2, ), dtype=torch.int32, device=sizes_i32.device, ) @@ -108,9 +107,7 @@ def _triton_via_route_a( ) -> Tuple[torch.Tensor, torch.Tensor]: from .triton_kernels.index import map_to_index as triton_map_to_index - mask_64, sizes_64 = _expand_mask_and_sizes_256_to_64( - logical_mask_256, logical_kv_sizes_256 - ) + mask_64, sizes_64 = _expand_mask_and_sizes_256_to_64(logical_mask_256, logical_kv_sizes_256) q2k_idx, q2k_num = triton_map_to_index(mask_64.to(torch.bool)) return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64) @@ -127,13 +124,9 @@ def block_sparse_attn_256( logical_block_map_256 = logical_block_map_256.unsqueeze(0) if _resolve_backend() == "triton": - return _triton_via_route_a( - q, k, v, logical_block_map_256, logical_variable_block_sizes_256 - ) + return _triton_via_route_a(q, k, v, logical_block_map_256, logical_variable_block_sizes_256) - mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128( - logical_block_map_256, logical_variable_block_sizes_256 - ) + mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128(logical_block_map_256, logical_variable_block_sizes_256) from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd return block_sparse_attn_cute_fwd(q, k, v, mask_128, sizes_128) @@ -163,8 +156,6 @@ def block_sparse_attn_256_bshd( ) return out_bhsd.transpose(1, 2).contiguous(), aux - mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128( - logical_block_map_256, logical_variable_block_sizes_256 - ) + mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128(logical_block_map_256, logical_variable_block_sizes_256) from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd return block_sparse_attn_cute_fwd_bshd(q, k, v, mask_128, sizes_128) diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py index e8e96334f..498e5ee6d 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py @@ -22,13 +22,11 @@ from typing import Tuple import torch -_FA4_IMPORT_HINT = ( - "VSA-256 CuTe fastpath requires a FlashAttention-4 CuTe build that " - "provides `flash_attn.cute` with block-sparsity support (plus " - "`nvidia-cutlass-dsl` and `quack-kernels`). This is an optional " - "dependency; the default VSA-256 path is Triton. Install the FA4 CuTe " - "build and set FASTVIDEO_VSA_CUTEDSL=1 to enable the CuTe fastpath." -) +_FA4_IMPORT_HINT = ("VSA-256 CuTe fastpath requires a FlashAttention-4 CuTe build that " + "provides `flash_attn.cute` with block-sparsity support (plus " + "`nvidia-cutlass-dsl` and `quack-kernels`). This is an optional " + "dependency; the default VSA-256 path is Triton. Install the FA4 CuTe " + "build and set FASTVIDEO_VSA_CUTEDSL=1 to enable the CuTe fastpath.") def _load_fa4_cute(): @@ -54,17 +52,14 @@ def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: if block_map.dim() == 3: block_map = block_map.unsqueeze(0) if block_map.dim() != 4: - raise ValueError( - f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), " - f"got shape={tuple(block_map.shape)}" - ) + raise ValueError(f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), " + f"got shape={tuple(block_map.shape)}") if block_map.dtype != torch.bool: block_map = block_map.to(torch.bool) if not block_map.is_cuda: raise RuntimeError("block_map must be a CUDA tensor.") from fastvideo_kernel.triton_kernels.index import ( - map_to_index as triton_map_to_index, - ) + map_to_index as triton_map_to_index, ) return triton_map_to_index(block_map) @@ -84,10 +79,8 @@ def _aggregate_q_block_map( ) -> torch.Tensor: factor = q_sparse_block_size // q_block_size if factor <= 0 or q_sparse_block_size % q_block_size != 0: - raise ValueError( - f"q_sparse_block_size must be a positive multiple of " - f"q_block_size ({q_block_size}), got {q_sparse_block_size}" - ) + raise ValueError(f"q_sparse_block_size must be a positive multiple of " + f"q_block_size ({q_block_size}), got {q_sparse_block_size}") bsz, nhead, q_blocks, kv_blocks = block_map.shape q_blocks_sparse = (q_blocks + factor - 1) // factor pad_q = q_blocks_sparse * factor - q_blocks @@ -164,9 +157,7 @@ def _cute_forward( q_block_size=q_block_size, ) kv_full = (variable_block_sizes == kv_block_size).view(1, 1, 1, -1) - kv_partial = ( - (variable_block_sizes > 0) & (variable_block_sizes < kv_block_size) - ).view(1, 1, 1, -1) + kv_partial = ((variable_block_sizes > 0) & (variable_block_sizes < kv_block_size)).view(1, 1, 1, -1) full_map = sparse_map & kv_full mask_map = sparse_map & kv_partial diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_varlen.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_varlen.py index 60197ec36..804fa026b 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_varlen.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_varlen.py @@ -41,9 +41,7 @@ def _scatter_to_padded( for actual in sizes: actual = min(actual, src_end - src_pos) if actual > 0: - dst[:, :, dst_pos:dst_pos + actual, :] = ( - src[src_pos:src_pos + actual].transpose(0, 1).unsqueeze(0) - ) + dst[:, :, dst_pos:dst_pos + actual, :] = (src[src_pos:src_pos + actual].transpose(0, 1).unsqueeze(0)) src_pos += actual dst_pos += block_size @@ -68,9 +66,7 @@ def _gather_from_padded( for actual in sizes: actual = min(actual, dst_end - dst_pos) if actual > 0: - dst[dst_pos:dst_pos + actual] = ( - src[0, :, src_pos:src_pos + actual, :].transpose(0, 1) - ) + dst[dst_pos:dst_pos + actual] = (src[0, :, src_pos:src_pos + actual, :].transpose(0, 1)) dst_pos += actual src_pos += block_size @@ -131,9 +127,7 @@ def block_sparse_attn_varlen( if q_variable_block_sizes_list is not None: q_vbs_resolved.append(q_variable_block_sizes_list[i]) else: - q_vbs_resolved.append( - torch.full((n_q_blocks,), block_size, dtype=torch.int32) - ) + q_vbs_resolved.append(torch.full((n_q_blocks, ), block_size, dtype=torch.int32)) total_padded_q = sum(padded_q_lens) total_padded_kv = sum(padded_kv_lens) @@ -146,16 +140,31 @@ def block_sparse_attn_varlen( kv_offset = 0 for i in range(num_seqs): _scatter_to_padded( - q, q_vbs_resolved[i], block_size, - q_packed, q_offset, cu_q[i], cu_q[i + 1], + q, + q_vbs_resolved[i], + block_size, + q_packed, + q_offset, + cu_q[i], + cu_q[i + 1], ) _scatter_to_padded( - k, variable_block_sizes_list[i], block_size, - k_packed, kv_offset, cu_kv[i], cu_kv[i + 1], + k, + variable_block_sizes_list[i], + block_size, + k_packed, + kv_offset, + cu_kv[i], + cu_kv[i + 1], ) _scatter_to_padded( - v, variable_block_sizes_list[i], block_size, - v_packed, kv_offset, cu_kv[i], cu_kv[i + 1], + v, + variable_block_sizes_list[i], + block_size, + v_packed, + kv_offset, + cu_kv[i], + cu_kv[i + 1], ) q_offset += padded_q_lens[i] kv_offset += padded_kv_lens[i] @@ -164,12 +173,19 @@ def block_sparse_attn_varlen( max_kv_per_q = max(t.shape[-1] for t in q2k_idx_list) global_q2k_idx = torch.zeros( - 1, num_heads, total_q_blocks, max_kv_per_q, - dtype=torch.int32, device=device, + 1, + num_heads, + total_q_blocks, + max_kv_per_q, + dtype=torch.int32, + device=device, ) global_q2k_num = torch.zeros( - 1, num_heads, total_q_blocks, - dtype=torch.int32, device=device, + 1, + num_heads, + total_q_blocks, + dtype=torch.int32, + device=device, ) global_vbs_parts = [] @@ -191,16 +207,25 @@ def block_sparse_attn_varlen( global_vbs = torch.cat(global_vbs_parts, dim=0).to(torch.int32).contiguous() out_packed, _ = block_sparse_attn_from_indices( - q_packed, k_packed, v_packed, - global_q2k_idx, global_q2k_num, global_vbs, + q_packed, + k_packed, + v_packed, + global_q2k_idx, + global_q2k_num, + global_vbs, ) out = torch.zeros(cu_q[-1], num_heads, head_dim, device=device, dtype=dtype) q_offset = 0 for i in range(num_seqs): _gather_from_padded( - out_packed, q_vbs_resolved[i], block_size, - out, q_offset, cu_q[i], cu_q[i + 1], + out_packed, + q_vbs_resolved[i], + block_size, + out, + q_offset, + cu_q[i], + cu_q[i + 1], ) q_offset += padded_q_lens[i] diff --git a/fastvideo-kernel/python/fastvideo_kernel/ops.py b/fastvideo-kernel/python/fastvideo_kernel/ops.py index 28bc87d71..c87f7171b 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/ops.py @@ -27,9 +27,7 @@ def sliding_tile_attention( ) -> torch.Tensor: # Check if the specific op is available if sta_fwd is None: - return sliding_tile_attention_triton( - q, k, v, window_size, text_length, has_text, seq_shape - ) + return sliding_tile_attention_triton(q, k, v, window_size, text_length, has_text, seq_shape) seq_length = q.shape[2] shape_map = {"30x48x80": 1, "36x48x48": 2, "18x48x80": 3} @@ -53,11 +51,7 @@ def sliding_tile_attention( k_h = k[:, head_idx:head_idx + 1].contiguous() v_h = v[:, head_idx:head_idx + 1].contiguous() o_h = torch.empty_like(q_h) - sta_fwd( - q_h, k_h, - v_h, o_h, - t, h, w, text_length, False, has_text, flag - ) + sta_fwd(q_h, k_h, v_h, o_h, t, h, w, text_length, False, has_text, flag) output[:, head_idx:head_idx + 1] = o_h if has_text: @@ -96,35 +90,27 @@ def video_sparse_attn( if k.shape[0] != batch or v.shape[0] != batch or k.shape[1] != heads or v.shape[1] != heads: raise ValueError("Expected q/k/v to have the same batch and head dimensions.") if v.shape[2] != kv_seq_len: - raise ValueError( - f"Expected k and v to have the same sequence length, got " - f"k.shape[2]={kv_seq_len}, v.shape[2]={v.shape[2]}" - ) + raise ValueError(f"Expected k and v to have the same sequence length, got " + f"k.shape[2]={kv_seq_len}, v.shape[2]={v.shape[2]}") if q_seq_len % block_elements != 0 or kv_seq_len % block_elements != 0: - raise ValueError( - f"q_seq_len and kv_seq_len must be divisible by block_elements={block_elements}, " - f"got q_seq_len={q_seq_len}, kv_seq_len={kv_seq_len}" - ) + raise ValueError(f"q_seq_len and kv_seq_len must be divisible by block_elements={block_elements}, " + f"got q_seq_len={q_seq_len}, kv_seq_len={kv_seq_len}") q_num_blocks = q_seq_len // block_elements kv_num_blocks = kv_seq_len // block_elements if variable_block_sizes.numel() != kv_num_blocks: - raise ValueError( - f"variable_block_sizes must have length kv_num_blocks={kv_num_blocks}, " - f"got {variable_block_sizes.numel()}" - ) + raise ValueError(f"variable_block_sizes must have length kv_num_blocks={kv_num_blocks}, " + f"got {variable_block_sizes.numel()}") if q_variable_block_sizes.numel() != q_num_blocks: - raise ValueError( - f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, " - f"got {q_variable_block_sizes.numel()}" - ) + raise ValueError(f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, " + f"got {q_variable_block_sizes.numel()}") # Compression branch (fused Triton: bf16 read → fp32 accumulate → div → bf16 write) q_c = fused_block_mean(q, q_variable_block_sizes, block_elements) k_c = fused_block_mean(k, variable_block_sizes, block_elements) v_c = fused_block_mean(v, variable_block_sizes, block_elements) - scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim ** 0.5) + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim**0.5) attn = torch.softmax(scores, dim=-1) out_c = torch.matmul(attn, v_c) out_c = out_c.view(batch, heads, q_num_blocks, 1, dim) @@ -163,61 +149,46 @@ def video_sparse_attn_bshd( block_size = (block_size, block_size, block_size) block_elements = block_size[0] * block_size[1] * block_size[2] if block_elements != 256: - raise ValueError( - "video_sparse_attn_bshd is only defined for block_elements=256 " - f"(got {block_elements}); use video_sparse_attn for the 64-block path." - ) + raise ValueError("video_sparse_attn_bshd is only defined for block_elements=256 " + f"(got {block_elements}); use video_sparse_attn for the 64-block path.") batch, q_seq_len, heads, dim = q.shape kv_seq_len = k.shape[1] if k.shape[0] != batch or v.shape[0] != batch or k.shape[2] != heads or v.shape[2] != heads: raise ValueError("Expected q/k/v to have the same batch and head dimensions.") if v.shape[1] != kv_seq_len: - raise ValueError( - f"Expected k and v to have the same sequence length, got " - f"k.shape[1]={kv_seq_len}, v.shape[1]={v.shape[1]}" - ) + raise ValueError(f"Expected k and v to have the same sequence length, got " + f"k.shape[1]={kv_seq_len}, v.shape[1]={v.shape[1]}") if q_seq_len % block_elements != 0 or kv_seq_len % block_elements != 0: - raise ValueError( - f"q_seq_len and kv_seq_len must be divisible by block_elements={block_elements}, " - f"got q_seq_len={q_seq_len}, kv_seq_len={kv_seq_len}" - ) + raise ValueError(f"q_seq_len and kv_seq_len must be divisible by block_elements={block_elements}, " + f"got q_seq_len={q_seq_len}, kv_seq_len={kv_seq_len}") q_num_blocks = q_seq_len // block_elements kv_num_blocks = kv_seq_len // block_elements if variable_block_sizes.numel() != kv_num_blocks: - raise ValueError( - f"variable_block_sizes must have length kv_num_blocks={kv_num_blocks}, " - f"got {variable_block_sizes.numel()}" - ) + raise ValueError(f"variable_block_sizes must have length kv_num_blocks={kv_num_blocks}, " + f"got {variable_block_sizes.numel()}") if q_variable_block_sizes.numel() != q_num_blocks: - raise ValueError( - f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, " - f"got {q_variable_block_sizes.numel()}" - ) + raise ValueError(f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, " + f"got {q_variable_block_sizes.numel()}") # Compression branch (BSHD-native: mean over the 256-token axis). token_idx = torch.arange(block_elements, device=q.device, dtype=torch.int32) - q_token_valid = (token_idx.view(1, -1) < q_variable_block_sizes.view(-1, 1)).view( - 1, q_num_blocks, block_elements, 1, 1 - ) - kv_token_valid = (token_idx.view(1, -1) < variable_block_sizes.view(-1, 1)).view( - 1, kv_num_blocks, block_elements, 1, 1 - ) + q_token_valid = (token_idx.view(1, -1) < q_variable_block_sizes.view(-1, + 1)).view(1, q_num_blocks, block_elements, 1, 1) + kv_token_valid = (token_idx.view(1, -1) < variable_block_sizes.view(-1, + 1)).view(1, kv_num_blocks, block_elements, 1, 1) q_c = q.view(batch, q_num_blocks, block_elements, heads, dim) k_c = k.view(batch, kv_num_blocks, block_elements, heads, dim) v_c = v.view(batch, kv_num_blocks, block_elements, heads, dim) - q_c = ((q_c.float() * q_token_valid).sum(dim=2) - / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype) - k_c = ((k_c.float() * kv_token_valid).sum(dim=2) - / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype) - v_c = ((v_c.float() * kv_token_valid).sum(dim=2) - / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype) + q_c = ((q_c.float() * q_token_valid).sum(dim=2) / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype) + k_c = ((k_c.float() * kv_token_valid).sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype) + v_c = ((v_c.float() * kv_token_valid).sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype) q_ch = q_c.permute(0, 2, 1, 3).contiguous() k_ch = k_c.permute(0, 2, 1, 3).contiguous() v_ch = v_c.permute(0, 2, 1, 3).contiguous() - scores = torch.matmul(q_ch, k_ch.transpose(-2, -1)) / (dim ** 0.5) + scores = torch.matmul(q_ch, k_ch.transpose(-2, -1)) / (dim**0.5) attn = torch.softmax(scores, dim=-1) out_c_ch = torch.matmul(attn, v_ch) out_c_blk = out_c_ch.permute(0, 2, 1, 3).contiguous() @@ -229,9 +200,7 @@ def video_sparse_attn_bshd( out = out_s out_view = out.view(batch, q_num_blocks, block_elements, heads, dim) if compress_attn_weight is not None: - gate_view = compress_attn_weight.view( - batch, q_num_blocks, block_elements, heads, dim - ) + gate_view = compress_attn_weight.view(batch, q_num_blocks, block_elements, heads, dim) out_view.add_(out_c_blk.unsqueeze(2) * gate_view) else: out_view.add_(out_c_blk.unsqueeze(2)) diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py index fc8f249b9..52ef4294c 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py @@ -59,16 +59,8 @@ def _use_sm100_optimized_qat( use_global_sf_p: bool, ) -> bool: """Return whether this call matches the validated SM100 fast path.""" - return ( - _sm100_optimization_enabled() - and is_sm100(device) - and head_dim == 128 - and not causal - and is_qat - and fake_quant_p - and not two_level_quant_p - and not use_global_sf_p - ) + return (_sm100_optimization_enabled() and is_sm100(device) and head_dim == 128 and not causal and is_qat + and fake_quant_p and not two_level_quant_p and not use_global_sf_p) def _select_sm100_forward_config(n_ctx_q: int, n_ctx_kv: int, mode: str): @@ -92,12 +84,27 @@ def _mul_alpha(acc, alpha, BM: tl.constexpr, BN: tl.constexpr): @triton.jit -def _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid, - desc_k, desc_v, - offset_y, dtype: tl.constexpr, start_m, qk_scale, - BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr, - STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr, - N_CTX: tl.constexpr, warp_specialize: tl.constexpr, IS_HOPPER: tl.constexpr, +def _attn_fwd_inner(acc, + high_prec_acc, + l_i, + m_i, + q, + q_valid, + desc_k, + desc_v, + offset_y, + dtype: tl.constexpr, + start_m, + qk_scale, + BLOCK_M: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_N: tl.constexpr, + STAGE: tl.constexpr, + offs_m: tl.constexpr, + offs_n: tl.constexpr, + N_CTX: tl.constexpr, + warp_specialize: tl.constexpr, + IS_HOPPER: tl.constexpr, IS_QAT: tl.constexpr, fake_quant_P: tl.constexpr = True, two_level_quant_P: tl.constexpr = False, @@ -135,16 +142,14 @@ def _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid, m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale) qk = qk * qk_scale - m_ij[:, None] p = tl.math.exp2(qk) - if IS_QAT: - p, high_prec_p = fake_quantize( - src_tensor=p, - valid_src_mask=tl.full(shape=p.shape, value=1.0, dtype=p.dtype) == 1.0, - BLOCK_SIZE_OUT_DIM=BLOCK_M, - BLOCK_SIZE_QUANT_DIM=BLOCK_N, - dst_dtype=dtype, - two_level_quant_P=two_level_quant_P, - use_global_sf=use_global_sf_P - ) + if IS_QAT: + p, high_prec_p = fake_quantize(src_tensor=p, + valid_src_mask=tl.full(shape=p.shape, value=1.0, dtype=p.dtype) == 1.0, + BLOCK_SIZE_OUT_DIM=BLOCK_M, + BLOCK_SIZE_QUANT_DIM=BLOCK_N, + dst_dtype=dtype, + two_level_quant_P=two_level_quant_P, + use_global_sf=use_global_sf_P) l_ij = tl.sum(high_prec_p, 1) else: l_ij = tl.sum(p, 1) @@ -213,11 +218,11 @@ def _host_descriptor_pre_hook(nargs): NUM_STAGES_OPTIONS = [2, 3, 4] configs = [ - triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w, pre_hook=_host_descriptor_pre_hook) - for BM in [64, 128] - for BN in [32, 64, 128] - for s in NUM_STAGES_OPTIONS - for w in [4, 8] + triton.Config({ + 'BLOCK_M': BM, + 'BLOCK_N': BN + }, num_stages=s, num_warps=w, pre_hook=_host_descriptor_pre_hook) for BM in [64, 128] for BN in [32, 64, 128] + for s in NUM_STAGES_OPTIONS for w in [4, 8] ] if "PYTEST_VERSION" in os.environ: # Use a single config in testing for reproducibility @@ -238,8 +243,7 @@ def prune_invalid_configs(configs, named_args, **kwargs): N_CTX_KV = kwargs["N_CTX_KV"] return [ conf for conf in configs - if conf.kwargs.get("BLOCK_M", 0) <= N_CTX_Q - and conf.kwargs.get("BLOCK_N", 0) <= N_CTX_KV + if conf.kwargs.get("BLOCK_M", 0) <= N_CTX_Q and conf.kwargs.get("BLOCK_N", 0) <= N_CTX_KV and conf.kwargs.get("BLOCK_N", 0) % conf.kwargs.get("BLOCK_M", 0) == 0 ] @@ -255,21 +259,31 @@ def _maybe_make_tensor_desc(desc_or_ptr, shape, strides, block_shape): # @triton.autotune(configs=list(filter(keep, configs)), key=["N_CTX_Q", "N_CTX_KV", "HEAD_DIM", "FP8_OUTPUT", "warp_specialize"], # prune_configs_by={'early_config_prune': prune_invalid_configs}, cache_results=True) @triton.jit -def _attn_fwd(sm_scale, M, - Z, H, desc_q, desc_k, desc_v, desc_o, desc_high_prec_o, N_CTX_Q, N_CTX_KV, - HEAD_DIM: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, - FP8_OUTPUT: tl.constexpr, - STAGE: tl.constexpr, - warp_specialize: tl.constexpr, - IS_HOPPER: tl.constexpr, - IS_QAT: tl.constexpr, - fake_quant_P: tl.constexpr = True, - two_level_quant_P: tl.constexpr = False, - use_global_sf_P: tl.constexpr = True, - JOIN_QAT_PV: tl.constexpr = False, - ): +def _attn_fwd( + sm_scale, + M, + Z, + H, + desc_q, + desc_k, + desc_v, + desc_o, + desc_high_prec_o, + N_CTX_Q, + N_CTX_KV, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + FP8_OUTPUT: tl.constexpr, + STAGE: tl.constexpr, + warp_specialize: tl.constexpr, + IS_HOPPER: tl.constexpr, + IS_QAT: tl.constexpr, + fake_quant_P: tl.constexpr = True, + two_level_quant_P: tl.constexpr = False, + use_global_sf_P: tl.constexpr = True, + JOIN_QAT_PV: tl.constexpr = False, +): dtype = tl.float8e5 if FP8_OUTPUT else tl.bfloat16 tl.static_assert(BLOCK_N <= HEAD_DIM) start_m = tl.program_id(0) @@ -279,27 +293,33 @@ def _attn_fwd(sm_scale, M, y_dim_q = Z * H * N_CTX_Q y_dim_kv = Z * H * N_CTX_KV - desc_q = _maybe_make_tensor_desc(desc_q, shape=[y_dim_q, HEAD_DIM], strides=[HEAD_DIM, 1], + desc_q = _maybe_make_tensor_desc(desc_q, + shape=[y_dim_q, HEAD_DIM], + strides=[HEAD_DIM, 1], block_shape=[BLOCK_M, HEAD_DIM]) if FP8_OUTPUT: - desc_v = _maybe_make_tensor_desc(desc_v, shape=[HEAD_DIM, y_dim_kv], strides=[N_CTX_KV, 1], + desc_v = _maybe_make_tensor_desc(desc_v, + shape=[HEAD_DIM, y_dim_kv], + strides=[N_CTX_KV, 1], block_shape=[HEAD_DIM, BLOCK_N]) else: - desc_v = _maybe_make_tensor_desc(desc_v, shape=[y_dim_kv, HEAD_DIM], strides=[HEAD_DIM, 1], + desc_v = _maybe_make_tensor_desc(desc_v, + shape=[y_dim_kv, HEAD_DIM], + strides=[HEAD_DIM, 1], block_shape=[BLOCK_N, HEAD_DIM]) - desc_k = _maybe_make_tensor_desc(desc_k, shape=[y_dim_kv, HEAD_DIM], strides=[HEAD_DIM, 1], + desc_k = _maybe_make_tensor_desc(desc_k, + shape=[y_dim_kv, HEAD_DIM], + strides=[HEAD_DIM, 1], block_shape=[BLOCK_N, HEAD_DIM]) - desc_o = _maybe_make_tensor_desc( - desc_o, shape=[Z * H, N_CTX_Q, HEAD_DIM], - strides=[N_CTX_Q * HEAD_DIM, HEAD_DIM, 1], - block_shape=[1, BLOCK_M, HEAD_DIM] - ) + desc_o = _maybe_make_tensor_desc(desc_o, + shape=[Z * H, N_CTX_Q, HEAD_DIM], + strides=[N_CTX_Q * HEAD_DIM, HEAD_DIM, 1], + block_shape=[1, BLOCK_M, HEAD_DIM]) if IS_QAT: - desc_high_prec_o = _maybe_make_tensor_desc( - desc_high_prec_o, shape=[Z * H, N_CTX_Q, HEAD_DIM], - strides=[N_CTX_Q * HEAD_DIM, HEAD_DIM, 1], - block_shape=[1, BLOCK_M, HEAD_DIM] - ) + desc_high_prec_o = _maybe_make_tensor_desc(desc_high_prec_o, + shape=[Z * H, N_CTX_Q, HEAD_DIM], + strides=[N_CTX_Q * HEAD_DIM, HEAD_DIM, 1], + block_shape=[1, BLOCK_M, HEAD_DIM]) offset_y_q = off_z * (N_CTX_Q * H) + off_h * N_CTX_Q # offset for query tensor offset_y_kv = off_z * (N_CTX_KV * H) + off_h * N_CTX_KV # offset for key/value tensors @@ -327,24 +347,18 @@ def _attn_fwd(sm_scale, M, # For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE # For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE if STAGE & 1: - acc, high_prec_acc, l_i, m_i = _attn_fwd_inner( - acc, high_prec_acc, l_i, m_i, q, q_valid, - desc_k, desc_v, - offset_y_kv, dtype, start_m, qk_scale, - BLOCK_M, HEAD_DIM, BLOCK_N, - 4 - STAGE, offs_m, offs_n, N_CTX_KV, - warp_specialize, IS_HOPPER, IS_QAT, fake_quant_P, two_level_quant_P, use_global_sf_P, JOIN_QAT_PV - ) + acc, high_prec_acc, l_i, m_i = _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid, desc_k, desc_v, + offset_y_kv, dtype, start_m, qk_scale, BLOCK_M, HEAD_DIM, + BLOCK_N, 4 - STAGE, offs_m, offs_n, N_CTX_KV, warp_specialize, + IS_HOPPER, IS_QAT, fake_quant_P, two_level_quant_P, + use_global_sf_P, JOIN_QAT_PV) # stage 2: on-band if STAGE & 2: - acc, high_prec_acc, l_i, m_i = _attn_fwd_inner( - acc, high_prec_acc, l_i, m_i, q, q_valid, - desc_k, desc_v, - offset_y_kv, dtype, start_m, qk_scale, - BLOCK_M, HEAD_DIM, BLOCK_N, - 2, offs_m, offs_n, N_CTX_KV, - warp_specialize, IS_HOPPER, IS_QAT, fake_quant_P, two_level_quant_P, use_global_sf_P, JOIN_QAT_PV - ) + acc, high_prec_acc, l_i, m_i = _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid, desc_k, desc_v, + offset_y_kv, dtype, start_m, qk_scale, BLOCK_M, HEAD_DIM, + BLOCK_N, 2, offs_m, offs_n, N_CTX_KV, warp_specialize, IS_HOPPER, + IS_QAT, fake_quant_P, two_level_quant_P, use_global_sf_P, + JOIN_QAT_PV) # epilogue m_i += tl.math.log2(l_i) acc = acc / l_i[:, None] @@ -404,18 +418,15 @@ def _attn_fwd_exact_m( @triton.jit -def _attn_bwd_preprocess(O, DO, - Delta, - Z, H, N_CTX, - BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr - ): +def _attn_bwd_preprocess(O, DO, Delta, Z, H, N_CTX, BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr): off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) off_hz = tl.program_id(1) off_n = tl.arange(0, HEAD_DIM) valid = off_m < N_CTX # load o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :], mask=valid[:, None]) - do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :], mask=valid[:, None]).to(tl.float32) + do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :], + mask=valid[:, None]).to(tl.float32) delta = tl.sum(o * do, axis=1) # write-back tl.store(Delta + off_hz * N_CTX + off_m, delta, mask=valid) @@ -423,24 +434,36 @@ def _attn_bwd_preprocess(O, DO, # The main inner-loop logic for computing dK and dV. @triton.jit -def _attn_bwd_dkdv(dk, dv, - Q, k, v, qk_scale, - DO, - M, D, Q_MEAN, - # shared by Q/K/V/DO. - stride_tok, stride_d, - H, N_CTX, BLOCK_M1: tl.constexpr, - BLOCK_N1: tl.constexpr, - HEAD_DIM: tl.constexpr, - # Filled in by the wrapper. - start_n, start_m, num_steps, - MASK: tl.constexpr, - IS_QAT: tl.constexpr, - two_level_quant_P: tl.constexpr = False, - fake_quant_P: tl.constexpr = True, - SMOOTH_Q: tl.constexpr = False, - use_global_sf_P: tl.constexpr = True, - warp_specialize: tl.constexpr = False): +def _attn_bwd_dkdv( + dk, + dv, + Q, + k, + v, + qk_scale, + DO, + M, + D, + Q_MEAN, + # shared by Q/K/V/DO. + stride_tok, + stride_d, + H, + N_CTX, + BLOCK_M1: tl.constexpr, + BLOCK_N1: tl.constexpr, + HEAD_DIM: tl.constexpr, + # Filled in by the wrapper. + start_n, + start_m, + num_steps, + MASK: tl.constexpr, + IS_QAT: tl.constexpr, + two_level_quant_P: tl.constexpr = False, + fake_quant_P: tl.constexpr = True, + SMOOTH_Q: tl.constexpr = False, + use_global_sf_P: tl.constexpr = True, + warp_specialize: tl.constexpr = False): offs_m = start_m + tl.arange(0, BLOCK_M1) offs_n = start_n + tl.arange(0, BLOCK_N1) offs_k = tl.arange(0, HEAD_DIM) @@ -471,15 +494,13 @@ def _attn_bwd_dkdv(dk, dv, # Compute dV. p_quant = p if IS_QAT and fake_quant_P: - p_quant, _ = fake_quantize( - src_tensor=p, - valid_src_mask=tl.full(shape=p.shape, value=1.0, dtype=p.dtype) == 1.0, - BLOCK_SIZE_OUT_DIM=BLOCK_M1, - BLOCK_SIZE_QUANT_DIM=BLOCK_N1, - dst_dtype=p.dtype, - two_level_quant_P=two_level_quant_P, - use_global_sf=use_global_sf_P - ) + p_quant, _ = fake_quantize(src_tensor=p, + valid_src_mask=tl.full(shape=p.shape, value=1.0, dtype=p.dtype) == 1.0, + BLOCK_SIZE_OUT_DIM=BLOCK_M1, + BLOCK_SIZE_QUANT_DIM=BLOCK_N1, + dst_dtype=p.dtype, + two_level_quant_P=two_level_quant_P, + use_global_sf=use_global_sf_P) dv += tl.dot(tl.trans(p_quant.to(tl.bfloat16)), do) # D (= delta) is pre-divided by ds_scale. Di = tl.load(D + offs_m, mask=q_valid) @@ -501,20 +522,31 @@ def _attn_bwd_dkdv(dk, dv, # The main inner-loop logic for computing dQ @triton.jit -def _attn_bwd_dq(dq, q, K, V, - do, m, D, qk_scale, - # shared by Q/K/V/DO. - stride_tok, stride_d, - H, N_CTX, - K_MEAN, - BLOCK_M2: tl.constexpr, - BLOCK_N2: tl.constexpr, - HEAD_DIM: tl.constexpr, - # Filled in by the wrapper. - start_m, start_n, num_steps, - MASK: tl.constexpr, - SMOOTH_K: tl.constexpr, - warp_specialize: tl.constexpr = False): +def _attn_bwd_dq( + dq, + q, + K, + V, + do, + m, + D, + qk_scale, + # shared by Q/K/V/DO. + stride_tok, + stride_d, + H, + N_CTX, + K_MEAN, + BLOCK_M2: tl.constexpr, + BLOCK_N2: tl.constexpr, + HEAD_DIM: tl.constexpr, + # Filled in by the wrapper. + start_m, + start_n, + num_steps, + MASK: tl.constexpr, + SMOOTH_K: tl.constexpr, + warp_specialize: tl.constexpr = False): offs_m = start_m + tl.arange(0, BLOCK_M2) offs_n = start_n + tl.arange(0, BLOCK_N2) offs_k = tl.arange(0, HEAD_DIM) @@ -571,12 +603,25 @@ def _compute_cross_attn_pointer_offsets(bhid, H, N_CTX_Q, stride_z_q, stride_z_k @triton.jit -def _attn_bwd_dq_cross(Q, K, V, sm_scale, - DO, DQ, - M, D, - stride_z_q, stride_z_kv, stride_h_q, - stride_h_kv, stride_tok_q, stride_tok_kv, stride_d_q, stride_d_kv, - H, N_CTX_Q, N_CTX_KV, +def _attn_bwd_dq_cross(Q, + K, + V, + sm_scale, + DO, + DQ, + M, + D, + stride_z_q, + stride_z_kv, + stride_h_q, + stride_h_kv, + stride_tok_q, + stride_tok_kv, + stride_d_q, + stride_d_kv, + H, + N_CTX_Q, + N_CTX_KV, K_MEAN, BLOCK_M2: tl.constexpr, BLOCK_N2: tl.constexpr, @@ -588,7 +633,8 @@ def _attn_bwd_dq_cross(Q, K, V, sm_scale, qk_scale = sm_scale * RCP_LN2 bhid = tl.program_id(2) - off_chz, adj_q, adj_kv = _compute_cross_attn_pointer_offsets(bhid, H, N_CTX_Q, stride_z_q, stride_z_kv, stride_h_q, stride_h_kv) + off_chz, adj_q, adj_kv = _compute_cross_attn_pointer_offsets(bhid, H, N_CTX_Q, stride_z_q, stride_z_kv, stride_h_q, + stride_h_kv) Q += adj_q K += adj_kv @@ -640,12 +686,27 @@ def _attn_bwd_dq_cross(Q, K, V, sm_scale, @triton.jit -def _attn_bwd_dkdv_cross(Q, K, V, sm_scale, - DO, DK, DV, - M, D, Q_MEAN, - stride_z_q, stride_z_kv, stride_h_q, stride_h_kv, - stride_tok_q, stride_tok_kv, stride_d_q, stride_d_kv, - H, N_CTX_Q, N_CTX_KV, +def _attn_bwd_dkdv_cross(Q, + K, + V, + sm_scale, + DO, + DK, + DV, + M, + D, + Q_MEAN, + stride_z_q, + stride_z_kv, + stride_h_q, + stride_h_kv, + stride_tok_q, + stride_tok_kv, + stride_d_q, + stride_d_kv, + H, + N_CTX_Q, + N_CTX_KV, BLOCK_M1: tl.constexpr, BLOCK_N1: tl.constexpr, HEAD_DIM: tl.constexpr, @@ -654,14 +715,14 @@ def _attn_bwd_dkdv_cross(Q, K, V, sm_scale, fake_quant_P: tl.constexpr = True, SMOOTH_Q: tl.constexpr = False, use_global_sf_P: tl.constexpr = True, - warp_specialize: tl.constexpr = False - ): + warp_specialize: tl.constexpr = False): # Apply scale AFTER dot product for better precision RCP_LN2: tl.constexpr = 1.4426950408889634 # = 1.0 / ln(2) qk_scale = sm_scale * RCP_LN2 bhid = tl.program_id(2) - off_chz, adj_q, adj_kv = _compute_cross_attn_pointer_offsets(bhid, H, N_CTX_Q, stride_z_q, stride_z_kv, stride_h_q, stride_h_kv) + off_chz, adj_q, adj_kv = _compute_cross_attn_pointer_offsets(bhid, H, N_CTX_Q, stride_z_q, stride_z_kv, stride_h_q, + stride_h_kv) Q += adj_q K += adj_kv @@ -703,15 +764,13 @@ def _attn_bwd_dkdv_cross(Q, K, V, sm_scale, p = tl.math.exp2(qk - m[:, None]) p_quant = p if IS_QAT and fake_quant_P: - p_quant, _ = fake_quantize( - src_tensor=p, - valid_src_mask=tl.full(shape=p.shape, value=1.0, dtype=p.dtype) == 1.0, - BLOCK_SIZE_OUT_DIM=BLOCK_M1, - BLOCK_SIZE_QUANT_DIM=BLOCK_N1, - dst_dtype=p.dtype, - two_level_quant_P=two_level_quant_P, - use_global_sf=use_global_sf_P - ) + p_quant, _ = fake_quantize(src_tensor=p, + valid_src_mask=tl.full(shape=p.shape, value=1.0, dtype=p.dtype) == 1.0, + BLOCK_SIZE_OUT_DIM=BLOCK_M1, + BLOCK_SIZE_QUANT_DIM=BLOCK_N1, + dst_dtype=p.dtype, + two_level_quant_P=two_level_quant_P, + use_global_sf=use_global_sf_P) dv += tl.dot(tl.trans(p_quant.to(tl.bfloat16)), do) dp = tl.dot(do, tl.trans(v_block)) @@ -733,28 +792,40 @@ def _attn_bwd_dkdv_cross(Q, K, V, sm_scale, @triton.jit -def _attn_bwd(Q, K, V, sm_scale, - DO, - DQ, DK, DV, - M, D, Q_MEAN, - # shared by Q/K/V/DO. - stride_z, stride_h, stride_tok, stride_d, - H, N_CTX, - K_MEAN, - BLOCK_M1: tl.constexpr, - BLOCK_N1: tl.constexpr, - BLOCK_M2: tl.constexpr, - BLOCK_N2: tl.constexpr, - BLK_SLICE_FACTOR: tl.constexpr, - HEAD_DIM: tl.constexpr, - CAUSAL: tl.constexpr, - IS_QAT: tl.constexpr, - SMOOTH_K: tl.constexpr, - two_level_quant_P: tl.constexpr = False, - fake_quant_P: tl.constexpr = True, - SMOOTH_Q: tl.constexpr = False, - use_global_sf_P: tl.constexpr = True, - warp_specialize: tl.constexpr = False): +def _attn_bwd( + Q, + K, + V, + sm_scale, + DO, + DQ, + DK, + DV, + M, + D, + Q_MEAN, + # shared by Q/K/V/DO. + stride_z, + stride_h, + stride_tok, + stride_d, + H, + N_CTX, + K_MEAN, + BLOCK_M1: tl.constexpr, + BLOCK_N1: tl.constexpr, + BLOCK_M2: tl.constexpr, + BLOCK_N2: tl.constexpr, + BLK_SLICE_FACTOR: tl.constexpr, + HEAD_DIM: tl.constexpr, + CAUSAL: tl.constexpr, + IS_QAT: tl.constexpr, + SMOOTH_K: tl.constexpr, + two_level_quant_P: tl.constexpr = False, + fake_quant_P: tl.constexpr = True, + SMOOTH_Q: tl.constexpr = False, + use_global_sf_P: tl.constexpr = True, + warp_specialize: tl.constexpr = False): bhid = tl.program_id(2) off_chz = (bhid * N_CTX).to(tl.int64) @@ -804,23 +875,33 @@ def _attn_bwd(Q, K, V, sm_scale, if CAUSAL: num_steps = BLOCK_N1 // MASK_BLOCK_M1 - dk, dv = _attn_bwd_dkdv( - dk, dv, - Q, k, v, qk_scale, - DO, - M, D, Q_MEAN, - stride_tok, stride_d, - H, N_CTX, - MASK_BLOCK_M1, BLOCK_N1, HEAD_DIM, - start_n, start_m, num_steps, - MASK=True, - IS_QAT=IS_QAT, - two_level_quant_P=two_level_quant_P, - fake_quant_P=fake_quant_P, - SMOOTH_Q=SMOOTH_Q, - use_global_sf_P=use_global_sf_P, - warp_specialize=warp_specialize - ) + dk, dv = _attn_bwd_dkdv(dk, + dv, + Q, + k, + v, + qk_scale, + DO, + M, + D, + Q_MEAN, + stride_tok, + stride_d, + H, + N_CTX, + MASK_BLOCK_M1, + BLOCK_N1, + HEAD_DIM, + start_n, + start_m, + num_steps, + MASK=True, + IS_QAT=IS_QAT, + two_level_quant_P=two_level_quant_P, + fake_quant_P=fake_quant_P, + SMOOTH_Q=SMOOTH_Q, + use_global_sf_P=use_global_sf_P, + warp_specialize=warp_specialize) start_m += num_steps * MASK_BLOCK_M1 num_steps = (N_CTX - start_m + BLOCK_M1 - 1) // BLOCK_M1 @@ -830,23 +911,33 @@ def _attn_bwd(Q, K, V, sm_scale, num_steps = (N_CTX + BLOCK_M1 - 1) // BLOCK_M1 # Compute dK and dV for non-masked blocks. - dk, dv = _attn_bwd_dkdv( - dk, dv, - Q, k, v, qk_scale, - DO, - M, D, Q_MEAN, - stride_tok, stride_d, - H, N_CTX, - BLOCK_M1, BLOCK_N1, HEAD_DIM, - start_n, start_m, num_steps, - MASK=False, - IS_QAT=IS_QAT, - two_level_quant_P=two_level_quant_P, - fake_quant_P=fake_quant_P, - SMOOTH_Q=SMOOTH_Q, - use_global_sf_P=use_global_sf_P, - warp_specialize=warp_specialize - ) + dk, dv = _attn_bwd_dkdv(dk, + dv, + Q, + k, + v, + qk_scale, + DO, + M, + D, + Q_MEAN, + stride_tok, + stride_d, + H, + N_CTX, + BLOCK_M1, + BLOCK_N1, + HEAD_DIM, + start_n, + start_m, + num_steps, + MASK=False, + IS_QAT=IS_QAT, + two_level_quant_P=two_level_quant_P, + fake_quant_P=fake_quant_P, + SMOOTH_Q=SMOOTH_Q, + use_global_sf_P=use_global_sf_P, + warp_specialize=warp_specialize) dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d tl.store(dv_ptrs, dv, mask=kv_valid[:, None]) @@ -879,13 +970,25 @@ def _attn_bwd(Q, K, V, sm_scale, if CAUSAL: num_steps = BLOCK_M2 // MASK_BLOCK_N2 dq = _attn_bwd_dq( - dq, q, K, V, - do, m, D, qk_scale, - stride_tok, stride_d, - H, N_CTX, + dq, + q, + K, + V, + do, + m, + D, + qk_scale, + stride_tok, + stride_d, + H, + N_CTX, K_MEAN, - BLOCK_M2, MASK_BLOCK_N2, HEAD_DIM, - start_m, end_n - num_steps * MASK_BLOCK_N2, num_steps, + BLOCK_M2, + MASK_BLOCK_N2, + HEAD_DIM, + start_m, + end_n - num_steps * MASK_BLOCK_N2, + num_steps, MASK=True, SMOOTH_K=SMOOTH_K, warp_specialize=warp_specialize, @@ -900,13 +1003,25 @@ def _attn_bwd(Q, K, V, sm_scale, start_n = 0 # stage 2 dq = _attn_bwd_dq( - dq, q, K, V, - do, m, D, qk_scale, - stride_tok, stride_d, - H, N_CTX, + dq, + q, + K, + V, + do, + m, + D, + qk_scale, + stride_tok, + stride_d, + H, + N_CTX, K_MEAN, - BLOCK_M2, BLOCK_N2, HEAD_DIM, - start_m, start_n, num_steps, + BLOCK_M2, + BLOCK_N2, + HEAD_DIM, + start_m, + start_n, + num_steps, MASK=False, SMOOTH_K=SMOOTH_K, warp_specialize=warp_specialize, @@ -957,22 +1072,16 @@ class _attention(torch.autograd.Function): N_CTX_Q = q.shape[2] # Query sequence length N_CTX_KV = k.shape[2] # Key/Value sequence length (may differ from query) assert k.shape[2] == v.shape[2], "k and v must have the same sequence length" - sm100_optimized = ( - q.dtype == torch.bfloat16 - and k.dtype == q.dtype - and v.dtype == q.dtype - and k.device == q.device - and v.device == q.device - and _use_sm100_optimized_qat( - q.device, - HEAD_DIM_K, - causal, - IS_QAT, - fake_quant_P, - two_level_quant_P, - use_global_sf_P, - ) - ) + sm100_optimized = (q.dtype == torch.bfloat16 and k.dtype == q.dtype and v.dtype == q.dtype + and k.device == q.device and v.device == q.device and _use_sm100_optimized_qat( + q.device, + HEAD_DIM_K, + causal, + IS_QAT, + fake_quant_P, + two_level_quant_P, + use_global_sf_P, + )) # smoothing k from SageAttn ctx.k_mean = None @@ -980,7 +1089,7 @@ class _attention(torch.autograd.Function): k_mean = k.mean(dim=(0, 1, 2), keepdim=True) k = k - k_mean ctx.k_mean = k_mean.view(-1) - + q_orig, k_orig, v_orig = None, None, None if not use_qat_qkv_backward: q_orig, k_orig, v_orig = q, k, v @@ -1006,30 +1115,28 @@ class _attention(torch.autograd.Function): # (Z, H, N_CTX_Q, HEAD_DIM_K) -> (Z*H*N_CTX_Q, HEAD_DIM_K) desc_q = TensorDescriptor(q, shape=[y_dim_q, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], block_shape=dummy_block) if q.dtype == torch.float8_e5m2: - desc_v = TensorDescriptor( - v, shape=[HEAD_DIM_K, y_dim_kv], strides=[k.shape[2], 1], - block_shape=dummy_block - ) + desc_v = TensorDescriptor(v, + shape=[HEAD_DIM_K, y_dim_kv], + strides=[k.shape[2], 1], + block_shape=dummy_block) else: - desc_v = TensorDescriptor( - v, shape=[y_dim_kv, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], - block_shape=dummy_block - ) + desc_v = TensorDescriptor(v, + shape=[y_dim_kv, HEAD_DIM_K], + strides=[HEAD_DIM_K, 1], + block_shape=dummy_block) desc_k = TensorDescriptor(k, shape=[y_dim_kv, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], block_shape=dummy_block) # Use 3D descriptor for output to handle N_CTX_Q boundaries correctly dummy_block3d = [1, 1, 1] ZH = q.shape[0] * q.shape[1] - desc_o = TensorDescriptor( - o, shape=[ZH, N_CTX_Q, HEAD_DIM_K], - strides=[N_CTX_Q * HEAD_DIM_K, HEAD_DIM_K, 1], - block_shape=dummy_block3d - ) + desc_o = TensorDescriptor(o, + shape=[ZH, N_CTX_Q, HEAD_DIM_K], + strides=[N_CTX_Q * HEAD_DIM_K, HEAD_DIM_K, 1], + block_shape=dummy_block3d) if IS_QAT: - desc_high_prec_o = TensorDescriptor( - high_prec_o, shape=[ZH, N_CTX_Q, HEAD_DIM_K], - strides=[N_CTX_Q * HEAD_DIM_K, HEAD_DIM_K, 1], - block_shape=dummy_block3d - ) + desc_high_prec_o = TensorDescriptor(high_prec_o, + shape=[ZH, N_CTX_Q, HEAD_DIM_K], + strides=[N_CTX_Q * HEAD_DIM_K, HEAD_DIM_K, 1], + block_shape=dummy_block3d) else: desc_high_prec_o = desc_o # Use regular output descriptor when not in QAT mode else: @@ -1065,13 +1172,10 @@ class _attention(torch.autograd.Function): if sm100_optimized: fwd_mode = os.environ.get("FASTVIDEO_ATTN_QAT_FWD_MODE", "fast").lower() if fwd_mode not in {"fast", "balanced", "reference"}: - raise ValueError( - f"FASTVIDEO_ATTN_QAT_FWD_MODE={fwd_mode!r} " - "(want fast|balanced|reference)" - ) + raise ValueError(f"FASTVIDEO_ATTN_QAT_FWD_MODE={fwd_mode!r} " + "(want fast|balanced|reference)") fwd_block_m, fwd_block_n, fwd_num_warps, fwd_num_stages = _select_sm100_forward_config( - N_CTX_Q, N_CTX_KV, fwd_mode - ) + N_CTX_Q, N_CTX_KV, fwd_mode) if IS_QAT: fake_q = torch.empty_like(q) fake_k = torch.empty_like(k) @@ -1080,9 +1184,18 @@ class _attention(torch.autograd.Function): # override desc_q, desc_k, desc_v with fake_q, fake_k, fake_v if supports_host_descriptor() and not (is_hopper() and warp_specialize): # (Z, H, N_CTX_Q, HEAD_DIM_K) -> (Z*H*N_CTX_Q, HEAD_DIM_K) - desc_q = TensorDescriptor(fake_q, shape=[y_dim_q, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], block_shape=dummy_block) - desc_k = TensorDescriptor(fake_k, shape=[y_dim_kv, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], block_shape=dummy_block) - desc_v = TensorDescriptor(fake_v, shape=[y_dim_kv, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], block_shape=dummy_block) + desc_q = TensorDescriptor(fake_q, + shape=[y_dim_q, HEAD_DIM_K], + strides=[HEAD_DIM_K, 1], + block_shape=dummy_block) + desc_k = TensorDescriptor(fake_k, + shape=[y_dim_kv, HEAD_DIM_K], + strides=[HEAD_DIM_K, 1], + block_shape=dummy_block) + desc_v = TensorDescriptor(fake_v, + shape=[y_dim_kv, HEAD_DIM_K], + strides=[HEAD_DIM_K, 1], + block_shape=dummy_block) else: desc_q = fake_q desc_k = fake_k @@ -1093,23 +1206,39 @@ class _attention(torch.autograd.Function): grid_2 = (triton.cdiv(k.shape[2], qkv_block_n), q.shape[0] * q.shape[1], 1) fake_quantize_q[grid_1]( - q, fake_q, - q.stride(0), q.stride(1), - q.stride(2), q.stride(3), - fake_q.stride(0), fake_q.stride(1), - fake_q.stride(2), fake_q.stride(3), - H, N_CTX_Q, - BLOCK_M=qkv_block_m, HEAD_DIM=HEAD_DIM_K, + q, + fake_q, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + fake_q.stride(0), + fake_q.stride(1), + fake_q.stride(2), + fake_q.stride(3), + H, + N_CTX_Q, + BLOCK_M=qkv_block_m, + HEAD_DIM=HEAD_DIM_K, use_global_sf=use_global_sf_QKV, ) fake_quantize_kv[grid_2]( - k, v, fake_k, fake_v, - k.stride(0), k.stride(1), - k.stride(2), k.stride(3), - fake_k.stride(0), fake_k.stride(1), - fake_k.stride(2), fake_k.stride(3), - H, N_CTX_KV, - BLOCK_N=qkv_block_n, HEAD_DIM=HEAD_DIM_K, + k, + v, + fake_k, + fake_v, + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + fake_k.stride(0), + fake_k.stride(1), + fake_k.stride(2), + fake_k.stride(3), + H, + N_CTX_KV, + BLOCK_N=qkv_block_n, + HEAD_DIM=HEAD_DIM_K, use_global_sf=use_global_sf_QKV, ) @@ -1126,35 +1255,35 @@ class _attention(torch.autograd.Function): "FP8_OUTPUT": q.dtype == torch.float8_e5m2, }) - _attn_fwd[grid]( - sm_scale, M, - q.shape[0], q.shape[1], - desc_q, desc_k, desc_v, desc_o, desc_high_prec_o, - N_CTX_Q=N_CTX_Q, - N_CTX_KV=N_CTX_KV, - HEAD_DIM=HEAD_DIM_K, - BLOCK_M=fwd_block_m, BLOCK_N=fwd_block_n, - FP8_OUTPUT=q.dtype == torch.float8_e5m2, - STAGE=stage, - warp_specialize=warp_specialize, - IS_HOPPER=is_hopper(), - IS_QAT=IS_QAT, - fake_quant_P=fake_quant_P, - two_level_quant_P=two_level_quant_P, - use_global_sf_P=use_global_sf_P, - JOIN_QAT_PV=(consumer_blackwell and _consumer_blackwell_join_qat_pv_enabled()), - num_warps=fwd_num_warps, - num_stages=fwd_num_stages, - **extra_kern_args - ) + _attn_fwd[grid](sm_scale, + M, + q.shape[0], + q.shape[1], + desc_q, + desc_k, + desc_v, + desc_o, + desc_high_prec_o, + N_CTX_Q=N_CTX_Q, + N_CTX_KV=N_CTX_KV, + HEAD_DIM=HEAD_DIM_K, + BLOCK_M=fwd_block_m, + BLOCK_N=fwd_block_n, + FP8_OUTPUT=q.dtype == torch.float8_e5m2, + STAGE=stage, + warp_specialize=warp_specialize, + IS_HOPPER=is_hopper(), + IS_QAT=IS_QAT, + fake_quant_P=fake_quant_P, + two_level_quant_P=two_level_quant_P, + use_global_sf_P=use_global_sf_P, + JOIN_QAT_PV=(consumer_blackwell and _consumer_blackwell_join_qat_pv_enabled()), + num_warps=fwd_num_warps, + num_stages=fwd_num_stages, + **extra_kern_args) exact_m = _sm100_exact_m_enabled() - if ( - sm100_optimized - and fwd_mode != "reference" - and exact_m - and (fwd_block_m != 32 or fwd_block_n != 32) - ): + if (sm100_optimized and fwd_mode != "reference" and exact_m and (fwd_block_m != 32 or fwd_block_n != 32)): # The large forward tile changes a legal reduction order. Restore # the legacy statistic so dV remains bitwise-compatible while the # two output paths retain the faster large-tile PV computation. @@ -1226,25 +1355,22 @@ class _attention(torch.autograd.Function): arg_k = k pre_grid = ((N_CTX_Q + PRE_BLOCK - 1) // PRE_BLOCK, BATCH * N_HEAD) delta = torch.empty_like(M) - _attn_bwd_preprocess[pre_grid]( - o_for_bwd, do, - delta, - BATCH, N_HEAD, N_CTX_Q, - BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM - ) + _attn_bwd_preprocess[pre_grid](o_for_bwd, + do, + delta, + BATCH, + N_HEAD, + N_CTX_Q, + BLOCK_M=PRE_BLOCK, + HEAD_DIM=ctx.HEAD_DIM) q_m = None if ctx.smooth_q: # _, q_m = triton_group_mean(q) q_m = q_m.repeat_interleave(q.shape[2] // q_m.shape[2], dim=2) # B,H,L,D - sm100_optimized_backward = ( - getattr(ctx, "sm100_optimized", False) - and ctx.use_qat_qkv_backward - and not ctx.smooth_k - and not ctx.smooth_q - and N_CTX_KV % 16 == 0 - ) + sm100_optimized_backward = (getattr(ctx, "sm100_optimized", False) and ctx.use_qat_qkv_backward + and not ctx.smooth_k and not ctx.smooth_q and N_CTX_KV % 16 == 0) if sm100_optimized_backward: # Keeping dQ and dK/dV in separate programs allows 64x64 tiles # without carrying all three fp32 accumulators at once. On SM100 @@ -1319,36 +1445,66 @@ class _attention(torch.autograd.Function): elif N_CTX_Q == N_CTX_KV: # Use existing kernel for self-attention (same sequence lengths) grid = ((N_CTX_KV + BLOCK_N1 - 1) // BLOCK_N1, 1, BATCH * N_HEAD) - _attn_bwd[grid]( - q, arg_k, v, ctx.sm_scale, do, dq, dk, dv, - M, delta, q_m, - q.stride(0), q.stride(1), q.stride(2), q.stride(3), - N_HEAD, N_CTX_KV, - ctx.k_mean, - BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1, - BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2, - BLK_SLICE_FACTOR=BLK_SLICE_FACTOR, - HEAD_DIM=ctx.HEAD_DIM, - CAUSAL=ctx.causal, - IS_QAT=ctx.IS_QAT, - SMOOTH_K=ctx.smooth_k, - two_level_quant_P=ctx.two_level_quant_P, - fake_quant_P=ctx.fake_quant_P, - SMOOTH_Q=ctx.smooth_q, - use_global_sf_P=ctx.use_global_sf_P, - warp_specialize=ctx.warp_specialize, - num_warps=NUM_WARPS, - num_stages=NUM_STAGES - ) + _attn_bwd[grid](q, + arg_k, + v, + ctx.sm_scale, + do, + dq, + dk, + dv, + M, + delta, + q_m, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + N_HEAD, + N_CTX_KV, + ctx.k_mean, + BLOCK_M1=BLOCK_M1, + BLOCK_N1=BLOCK_N1, + BLOCK_M2=BLOCK_M2, + BLOCK_N2=BLOCK_N2, + BLK_SLICE_FACTOR=BLK_SLICE_FACTOR, + HEAD_DIM=ctx.HEAD_DIM, + CAUSAL=ctx.causal, + IS_QAT=ctx.IS_QAT, + SMOOTH_K=ctx.smooth_k, + two_level_quant_P=ctx.two_level_quant_P, + fake_quant_P=ctx.fake_quant_P, + SMOOTH_Q=ctx.smooth_q, + use_global_sf_P=ctx.use_global_sf_P, + warp_specialize=ctx.warp_specialize, + num_warps=NUM_WARPS, + num_stages=NUM_STAGES) else: # Use separate kernels for cross-attention (different sequence lengths) grid_dq = ((N_CTX_Q + BLOCK_M2 - 1) // BLOCK_M2, 1, BATCH * N_HEAD) _attn_bwd_dq_cross[grid_dq]( - q, arg_k, v, ctx.sm_scale, do, dq, M, delta, - q.stride(0), k.stride(0), q.stride(1), k.stride(1), q.stride(2), k.stride(2), q.stride(3), k.stride(3), - N_HEAD, N_CTX_Q, N_CTX_KV, + q, + arg_k, + v, + ctx.sm_scale, + do, + dq, + M, + delta, + q.stride(0), + k.stride(0), + q.stride(1), + k.stride(1), + q.stride(2), + k.stride(2), + q.stride(3), + k.stride(3), + N_HEAD, + N_CTX_Q, + N_CTX_KV, ctx.k_mean, - BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2, + BLOCK_M2=BLOCK_M2, + BLOCK_N2=BLOCK_N2, HEAD_DIM=ctx.HEAD_DIM, SMOOTH_K=ctx.smooth_k, warp_specialize=ctx.warp_specialize, @@ -1357,10 +1513,29 @@ class _attention(torch.autograd.Function): ) grid_dkdv = ((N_CTX_KV + BLOCK_N1 - 1) // BLOCK_N1, 1, BATCH * N_HEAD) _attn_bwd_dkdv_cross[grid_dkdv]( - q, arg_k, v, ctx.sm_scale, do, dk, dv, M, delta, q_m, - q.stride(0), k.stride(0), q.stride(1), k.stride(1), q.stride(2), k.stride(2), q.stride(3), k.stride(3), - N_HEAD, N_CTX_Q, N_CTX_KV, - BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1, + q, + arg_k, + v, + ctx.sm_scale, + do, + dk, + dv, + M, + delta, + q_m, + q.stride(0), + k.stride(0), + q.stride(1), + k.stride(1), + q.stride(2), + k.stride(2), + q.stride(3), + k.stride(3), + N_HEAD, + N_CTX_Q, + N_CTX_KV, + BLOCK_M1=BLOCK_M1, + BLOCK_N1=BLOCK_N1, HEAD_DIM=ctx.HEAD_DIM, IS_QAT=ctx.IS_QAT, two_level_quant_P=ctx.two_level_quant_P, diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py index e0ac57db7..71159b43b 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py @@ -104,8 +104,7 @@ def _attn_fwd_sparse( block_shape=(HEAD_DIM, BLOCK_N), order=(0, 1)) - v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, - 0) + v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0) V_base = tl.make_block_ptr(base=V + v_off, shape=(N_CTX_KV, HEAD_DIM), strides=(stride_vk, stride_vn), @@ -178,10 +177,8 @@ def _attn_bwd_preprocess( off_hz = tl.program_id(1) off_n = tl.arange(0, HEAD_DIM) # load - o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + - off_n[None, :]) - do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + - off_n[None, :]).to(tl.float32) + o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]) + do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32) delta = tl.sum(o * do, axis=1) # write-back tl.store(Delta + off_hz * N_CTX + off_m, delta) @@ -235,8 +232,7 @@ def _attn_bwd_dkdv( block_size = tl.load(variable_block_sizes + kv_blk) for blk_idx in range(q_blocks * 2): - block_sparse_offset = (tl.load(q_ptr + blk_idx // 2).to(tl.int32) * 2 + - blk_idx % 2) * step_m + block_sparse_offset = (tl.load(q_ptr + blk_idx // 2).to(tl.int32) * 2 + blk_idx % 2) * step_m qT = tl.load(qT_ptrs + block_sparse_offset * stride_tok) # Load m before computing qk to reduce pipeline stall. offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1) @@ -497,7 +493,7 @@ def _attn_bwd_dkdv_kernel( stride_tok, stride_d, # # batch/head strides (may differ between Q and KV) - stride_qz, + stride_qz, stride_qh, stride_kz, stride_kh, @@ -607,7 +603,7 @@ def _attn_bwd_dq_kernel( stride_tok, stride_d, # # batch/head strides (may differ between Q and KV) - stride_qz, + stride_qz, stride_qh, stride_kz, stride_kh, @@ -689,8 +685,7 @@ def _attn_bwd_dq_kernel( # ──────────────────────────── SPARSE ADDITION BEGIN ─────────────────────────── -def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, - variable_block_sizes): +def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, variable_block_sizes): B, H, Tq, D = q.shape Tkv = k.shape[2] sm_scale = 1.0 / math.sqrt(D) @@ -701,8 +696,7 @@ def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, -1] == Tq // 64, f"shape mismatch, Tq // 64 = {Tq // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}" assert variable_block_sizes.numel() == Tkv // 64, ( f"shape mismatch, variable_block_sizes must have length {Tkv // 64}, " - f"got {variable_block_sizes.numel()}" - ) + f"got {variable_block_sizes.numel()}") o = torch.empty_like(q) M = torch.empty((B, H, Tq), dtype=torch.float32, device=q.device) @@ -743,8 +737,7 @@ def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, return o, M -def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, - k2q_index, k2q_num, variable_block_sizes): +def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num, variable_block_sizes): assert do.is_contiguous() B, H, Tq, D = q.shape diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/fused_compress_topk.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/fused_compress_topk.py index 68a4acd04..c450f192e 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/fused_compress_topk.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/fused_compress_topk.py @@ -122,9 +122,13 @@ def _fused_block_mean_bwd( grid = (num_blocks, B * H) _fused_block_mean_bwd_kernel[grid]( - go_flat, gx_flat, variable_block_sizes, - go_flat.stride(0), go_flat.stride(1), - gx_flat.stride(0), gx_flat.stride(1), + go_flat, + gx_flat, + variable_block_sizes, + go_flat.stride(0), + go_flat.stride(1), + gx_flat.stride(0), + gx_flat.stride(1), num_blocks, BLOCK_ELEMENTS=block_elements, HEAD_DIM=D, @@ -152,9 +156,13 @@ def _fused_block_mean_fwd( grid = (num_blocks, B * H) _fused_block_mean_kernel[grid]( - x_flat, out_flat, variable_block_sizes, - x_flat.stride(0), x_flat.stride(1), - out_flat.stride(0), out_flat.stride(1), + x_flat, + out_flat, + variable_block_sizes, + x_flat.stride(0), + x_flat.stride(1), + out_flat.stride(0), + out_flat.stride(1), num_blocks, BLOCK_ELEMENTS=block_elements, HEAD_DIM=D, @@ -311,7 +319,8 @@ def fused_topk_mask( logger.debug( "fused_topk_mask: kv_blocks=%d exceeds Triton limit %d, " "falling back to PyTorch topk (slower)", - kv_blocks, MAX_KV_BLOCK_SIZE, + kv_blocks, + MAX_KV_BLOCK_SIZE, ) return _pytorch_topk_mask_fallback(scores, topk) @@ -323,9 +332,14 @@ def fused_topk_mask( grid = (q_blocks, B * H) _fused_topk_mask_kernel[grid]( - scores_flat, mask_flat, - scores_flat.stride(0), scores_flat.stride(1), scores_flat.stride(2), - mask_flat.stride(0), mask_flat.stride(1), mask_flat.stride(2), + scores_flat, + mask_flat, + scores_flat.stride(0), + scores_flat.stride(1), + scores_flat.stride(2), + mask_flat.stride(0), + mask_flat.stride(1), + mask_flat.stride(2), kv_blocks=kv_blocks, topk=topk, KV_BLOCK_SIZE=KV_BLOCK_SIZE, diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py index 65e479331..be7c5c379 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py @@ -5,6 +5,7 @@ import triton import triton.language as tl import torch + @triton.jit def topk_index_to_map_kernel( map_ptr, @@ -57,14 +58,10 @@ def map_to_index_kernel( tl.store(index_ptr_base + num * index_kv_stride, i) num += 1 - tl.store( - index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride + - q * index_num_q_stride, num) + tl.store(index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride + q * index_num_q_stride, num) -def topk_index_to_map(index: torch.Tensor, - num_kv_blocks: int, - transpose_map: bool = False): +def topk_index_to_map(index: torch.Tensor, num_kv_blocks: int, transpose_map: bool = False): """ Convert topk indices to a map. @@ -83,13 +80,9 @@ def topk_index_to_map(index: torch.Tensor, bs, h, num_q_blocks, topk = index.shape if transpose_map is False: - block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks), - dtype=torch.bool, - device=index.device) + block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks), dtype=torch.bool, device=index.device) else: - block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks), - dtype=torch.bool, - device=index.device) + block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks), dtype=torch.bool, device=index.device) block_map = block_map.transpose(2, 3) grid = (bs, h, num_q_blocks) @@ -126,13 +119,8 @@ def map_to_index(block_map: torch.Tensor): """ bs, h, num_q_blocks, num_kv_blocks = block_map.shape - index = torch.full((block_map.shape), - -1, - dtype=torch.int32, - device=block_map.device) - index_num = torch.empty((bs, h, num_q_blocks), - dtype=torch.int32, - device=block_map.device) + index = torch.full((block_map.shape), -1, dtype=torch.int32, device=block_map.device) + index_num = torch.empty((bs, h, num_q_blocks), dtype=torch.int32, device=block_map.device) grid = (bs, h, num_q_blocks) map_to_index_kernel[grid]( @@ -162,10 +150,20 @@ def _invert_indices_kernel( q2k_num_ptr, k2q_idx_ptr, k2q_num_ptr, - q2k_idx_b, q2k_idx_h, q2k_idx_q, q2k_idx_k, - q2k_num_b, q2k_num_h, q2k_num_q, - k2q_idx_b, k2q_idx_h, k2q_idx_k, k2q_idx_q, - k2q_num_b, k2q_num_h, k2q_num_k, + q2k_idx_b, + q2k_idx_h, + q2k_idx_q, + q2k_idx_k, + q2k_num_b, + q2k_num_h, + q2k_num_q, + k2q_idx_b, + k2q_idx_h, + k2q_idx_k, + k2q_idx_q, + k2q_num_b, + k2q_num_h, + k2q_num_k, MAX_KV_PER_Q: tl.constexpr, ): # One program per (b, h, q): reserve a slot in k2q via atomicAdd, write q. @@ -173,36 +171,17 @@ def _invert_indices_kernel( pid_h = tl.program_id(1) pid_q = tl.program_id(2) - n = tl.load( - q2k_num_ptr - + pid_b * q2k_num_b - + pid_h * q2k_num_h - + pid_q * q2k_num_q - ) + n = tl.load(q2k_num_ptr + pid_b * q2k_num_b + pid_h * q2k_num_h + pid_q * q2k_num_q) - q2k_row = ( - q2k_idx_ptr - + pid_b * q2k_idx_b - + pid_h * q2k_idx_h - + pid_q * q2k_idx_q - ) + q2k_row = (q2k_idx_ptr + pid_b * q2k_idx_b + pid_h * q2k_idx_h + pid_q * q2k_idx_q) for i in tl.range(0, MAX_KV_PER_Q): if i < n: kv = tl.load(q2k_row + i * q2k_idx_k) - count_ptr = ( - k2q_num_ptr - + pid_b * k2q_num_b - + pid_h * k2q_num_h - + kv * k2q_num_k - ) + count_ptr = (k2q_num_ptr + pid_b * k2q_num_b + pid_h * k2q_num_h + kv * k2q_num_k) pos = tl.atomic_add(count_ptr, 1) tl.store( - k2q_idx_ptr - + pid_b * k2q_idx_b - + pid_h * k2q_idx_h - + kv * k2q_idx_k - + pos * k2q_idx_q, + k2q_idx_ptr + pid_b * k2q_idx_b + pid_h * k2q_idx_h + kv * k2q_idx_k + pos * k2q_idx_q, pid_q, ) @@ -214,22 +193,16 @@ def invert_indices( ) -> Tuple[torch.Tensor, torch.Tensor]: """Transpose a Q->KV index list into a K->Q one via atomic compaction (GPU).""" if q2k_idx.dim() != 4: - raise ValueError( - f"q2k_idx must be [B, H, Nq, Mk], got shape={tuple(q2k_idx.shape)}" - ) + raise ValueError(f"q2k_idx must be [B, H, Nq, Mk], got shape={tuple(q2k_idx.shape)}") if q2k_num.dim() != 3: - raise ValueError( - f"q2k_num must be [B, H, Nq], got shape={tuple(q2k_num.shape)}" - ) + raise ValueError(f"q2k_num must be [B, H, Nq], got shape={tuple(q2k_num.shape)}") if not q2k_idx.is_cuda or not q2k_num.is_cuda: raise RuntimeError("invert_indices requires CUDA tensors.") B, H, Nq, Mk = q2k_idx.shape if q2k_num.shape != (B, H, Nq): - raise ValueError( - f"q2k_num shape {tuple(q2k_num.shape)} does not match q2k_idx " - f"[B, H, Nq] = {(B, H, Nq)}" - ) + raise ValueError(f"q2k_num shape {tuple(q2k_num.shape)} does not match q2k_idx " + f"[B, H, Nq] = {(B, H, Nq)}") q2k_idx = q2k_idx.contiguous() q2k_num = q2k_num.contiguous() @@ -257,10 +230,20 @@ def invert_indices( q2k_num, k2q_idx, k2q_num, - q2k_idx.stride(0), q2k_idx.stride(1), q2k_idx.stride(2), q2k_idx.stride(3), - q2k_num.stride(0), q2k_num.stride(1), q2k_num.stride(2), - k2q_idx.stride(0), k2q_idx.stride(1), k2q_idx.stride(2), k2q_idx.stride(3), - k2q_num.stride(0), k2q_num.stride(1), k2q_num.stride(2), + q2k_idx.stride(0), + q2k_idx.stride(1), + q2k_idx.stride(2), + q2k_idx.stride(3), + q2k_num.stride(0), + q2k_num.stride(1), + q2k_num.stride(2), + k2q_idx.stride(0), + k2q_idx.stride(1), + k2q_idx.stride(2), + k2q_idx.stride(3), + k2q_num.stride(0), + k2q_num.stride(1), + k2q_num.stride(2), MAX_KV_PER_Q=Mk, ) diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py index 6fdbb36e9..7abc9f79a 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py @@ -8,6 +8,7 @@ from triton.language.target_info import cuda_capability_geq MXFP_BLOCK_SIZE = tl.constexpr(16) + @triton.jit def _compute_quant_and_scale( src_tensor, @@ -22,9 +23,7 @@ def _compute_quant_and_scale( is_fp4: tl.constexpr = mx_tensor_dtype == tl.uint8 tl.static_assert( - is_fp4 - or mx_tensor_dtype == tl.float8e4nv - or mx_tensor_dtype == tl.float8e5, + is_fp4 or mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5, "mx_tensor_dtype must be uint8, float8e4nv, or float8e5", ) @@ -32,7 +31,7 @@ def _compute_quant_and_scale( f32_tensor = src_tensor.to(tl.float32) abs_tensor = tl.abs(f32_tensor) abs_tensor = tl.where(valid_src_mask, abs_tensor, -1.0) # Don't consider padding tensors in scale computation - + if two_level_quant_P: # row max from SageAttn3 paper global_max_val = tl.max(f32_tensor, axis=1, keep_dims=True) # (BLOCK_SIZE_OUT_DIM, 1) @@ -51,8 +50,9 @@ def _compute_quant_and_scale( elif not two_level_quant_P and not use_global_sf: s_dec = 1.0 s_enc = 1.0 - - max_val = tl.max(abs_tensor, axis=2, keep_dims=True) # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1) # per block maxima + + max_val = tl.max(abs_tensor, axis=2, + keep_dims=True) # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1) # per block maxima s_dec_b = max_val / 6 # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1) s_dec_b_e4m3 = (s_dec_b * s_enc).to(tl.float8e4nv) # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1) s_enc_b = 1 / (s_dec_b_e4m3.to(tl.float32) * s_dec) # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1) @@ -132,6 +132,7 @@ def _compute_quant_and_scale( return out_tensor, dequant_scale, s_dec + @triton.jit def _compute_dequant( mx_tensor, @@ -141,7 +142,8 @@ def _compute_dequant( BLOCK_SIZE_QUANT_DIM: tl.constexpr, dst_dtype: tl.constexpr, ): - tl.static_assert(BLOCK_SIZE_QUANT_DIM % MXFP_BLOCK_SIZE == 0, f"Block size along quantization block must be a multiple of {MXFP_BLOCK_SIZE=}") + tl.static_assert(BLOCK_SIZE_QUANT_DIM % MXFP_BLOCK_SIZE == 0, + f"Block size along quantization block must be a multiple of {MXFP_BLOCK_SIZE=}") # uint8 signifies two fp4 e2m1 values packed into a single byte mx_tensor_dtype: tl.constexpr = mx_tensor.dtype tl.static_assert(dst_dtype == tl.float16 or dst_dtype == tl.bfloat16 or dst_dtype == tl.float32) diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py index 31d4acc79..bac2a75cf 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py @@ -3,33 +3,44 @@ import triton.language as tl from .nvfp4_utils import _compute_quant_and_scale, _compute_dequant + @triton.jit -def fake_quantize(src_tensor, valid_src_mask, BLOCK_SIZE_OUT_DIM: tl.constexpr, - BLOCK_SIZE_QUANT_DIM: tl.constexpr, - dst_dtype: tl.constexpr, - mx_tensor_dtype: tl.constexpr = tl.uint8, - use_global_sf: tl.constexpr = True, - two_level_quant_P: tl.constexpr = False): +def fake_quantize(src_tensor, + valid_src_mask, + BLOCK_SIZE_OUT_DIM: tl.constexpr, + BLOCK_SIZE_QUANT_DIM: tl.constexpr, + dst_dtype: tl.constexpr, + mx_tensor_dtype: tl.constexpr = tl.uint8, + use_global_sf: tl.constexpr = True, + two_level_quant_P: tl.constexpr = False): high_prec_src_tensor = src_tensor src_tensor, src_scale, src_s_dec = _compute_quant_and_scale(src_tensor=src_tensor, valid_src_mask=valid_src_mask, mx_tensor_dtype=mx_tensor_dtype, use_global_sf=use_global_sf, two_level_quant_P=two_level_quant_P) - src_tensor = _compute_dequant(mx_tensor=src_tensor, - scale=src_scale, - s_dec=src_s_dec, - BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM, - BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM, + src_tensor = _compute_dequant(mx_tensor=src_tensor, + scale=src_scale, + s_dec=src_s_dec, + BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM, + BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM, dst_dtype=dst_dtype) return src_tensor, high_prec_src_tensor.to(src_tensor.dtype) + @triton.jit -def fake_quantize_q(Q, fake_Q, stride_z_q, stride_h_q, - stride_tok_q, stride_d_q, - fake_stride_z_q, fake_stride_h_q, - fake_stride_tok_q, fake_stride_d_q, - H, N_CTX_Q, +def fake_quantize_q(Q, + fake_Q, + stride_z_q, + stride_h_q, + stride_tok_q, + stride_d_q, + fake_stride_z_q, + fake_stride_h_q, + fake_stride_tok_q, + fake_stride_d_q, + H, + N_CTX_Q, BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, use_global_sf: tl.constexpr = True): @@ -46,18 +57,33 @@ def fake_quantize_q(Q, fake_Q, stride_z_q, stride_h_q, q_valid = offs_m < N_CTX_Q q = tl.load(Q + offs_m[:, None] * stride_tok_q + offs_k[None, :] * stride_d_q, mask=q_valid[:, None], other=0.0) - q, _ = fake_quantize(src_tensor=q, valid_src_mask=q_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_M, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=q.dtype, use_global_sf=use_global_sf) + q, _ = fake_quantize(src_tensor=q, + valid_src_mask=q_valid[:, None], + BLOCK_SIZE_OUT_DIM=BLOCK_M, + BLOCK_SIZE_QUANT_DIM=HEAD_DIM, + dst_dtype=q.dtype, + use_global_sf=use_global_sf) tl.store(fake_Q + offs_m[:, None] * fake_stride_tok_q + offs_k[None, :] * fake_stride_d_q, q, mask=q_valid[:, None]) + @triton.jit -def fake_quantize_kv(K, V, fake_K, fake_V, stride_z_kv, stride_h_kv, - stride_tok_kv, stride_d_kv, - fake_stride_z_kv, fake_stride_h_kv, - fake_stride_tok_kv, fake_stride_d_kv, - H, N_CTX_KV, - BLOCK_N: tl.constexpr, - HEAD_DIM: tl.constexpr, - use_global_sf: tl.constexpr = True): +def fake_quantize_kv(K, + V, + fake_K, + fake_V, + stride_z_kv, + stride_h_kv, + stride_tok_kv, + stride_d_kv, + fake_stride_z_kv, + fake_stride_h_kv, + fake_stride_tok_kv, + fake_stride_d_kv, + H, + N_CTX_KV, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, + use_global_sf: tl.constexpr = True): bhid = tl.program_id(1) adj_kv = (stride_h_kv * (bhid % H) + stride_z_kv * (bhid // H)) fake_adj_kv = (fake_stride_h_kv * (bhid % H) + fake_stride_z_kv * (bhid // H)) @@ -70,11 +96,29 @@ def fake_quantize_kv(K, V, fake_K, fake_V, stride_z_kv, stride_h_kv, start_n = pid * BLOCK_N offs_n = start_n + tl.arange(0, BLOCK_N) offs_k = tl.arange(0, HEAD_DIM) - + kv_valid = offs_n < N_CTX_KV - k_block = tl.load(K + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None], other=0.0) - v_block = tl.load(V + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None], other=0.0) - k, _ = fake_quantize(src_tensor=k_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=k_block.dtype, use_global_sf=use_global_sf) - v, _ = fake_quantize(src_tensor=v_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=v_block.dtype, use_global_sf=use_global_sf) - tl.store(fake_K + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, k, mask=kv_valid[:, None]) - tl.store(fake_V + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, v, mask=kv_valid[:, None]) + k_block = tl.load(K + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, + mask=kv_valid[:, None], + other=0.0) + v_block = tl.load(V + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, + mask=kv_valid[:, None], + other=0.0) + k, _ = fake_quantize(src_tensor=k_block, + valid_src_mask=kv_valid[:, None], + BLOCK_SIZE_OUT_DIM=BLOCK_N, + BLOCK_SIZE_QUANT_DIM=HEAD_DIM, + dst_dtype=k_block.dtype, + use_global_sf=use_global_sf) + v, _ = fake_quantize(src_tensor=v_block, + valid_src_mask=kv_valid[:, None], + BLOCK_SIZE_OUT_DIM=BLOCK_N, + BLOCK_SIZE_QUANT_DIM=HEAD_DIM, + dst_dtype=v_block.dtype, + use_global_sf=use_global_sf) + tl.store(fake_K + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, + k, + mask=kv_valid[:, None]) + tl.store(fake_V + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, + v, + mask=kv_valid[:, None]) diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/sla_triton.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/sla_triton.py index 64f2f6dd5..945ee2e7a 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/sla_triton.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/sla_triton.py @@ -19,10 +19,14 @@ import triton.language as tl @triton.jit def _attn_fwd( - Q, K, V, + Q, + K, + V, qk_scale: tl.constexpr, topk: tl.constexpr, - LUT, LSE, OS, + LUT, + LSE, + OS, L: tl.constexpr, M_BLOCKS: tl.constexpr, D: tl.constexpr, @@ -45,7 +49,7 @@ def _attn_fwd( OS_ptrs = OS + qkv_offset + offs_m[:, None] * D + offs_d[None, :] LUT_ptr = LUT + lut_offset LSE_ptrs = LSE + lse_offset + offs_m - + m_i = tl.full([BLOCK_M], -float('inf'), dtype=tl.float32) l_i = tl.zeros([BLOCK_M], dtype=tl.float32) o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32) @@ -54,7 +58,7 @@ def _attn_fwd( for block_idx in tl.range(topk): idx_n = tl.load(LUT_ptr + block_idx) n_mask = offs_n < L - idx_n * BLOCK_N - + k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[None, :]) qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634) # = 1 / ln(2) if L - idx_n * BLOCK_N < BLOCK_N: @@ -76,14 +80,16 @@ def _attn_fwd( o_s = o_s / l_i[:, None] tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L) - + m_i += tl.math.log2(l_i) tl.store(LSE_ptrs, m_i, mask=offs_m < L) @triton.jit def _attn_bwd_preprocess( - OS, DOS, DELTAS, + OS, + DOS, + DELTAS, L, D: tl.constexpr, BLOCK_M: tl.constexpr, @@ -100,15 +106,21 @@ def _attn_bwd_preprocess( o_s = tl.load(OS + offs_m[:, None] * D + offs_d[None, :], mask=offs_m[:, None] < L) do_s = tl.load(DOS + offs_m[:, None] * D + offs_d[None, :], mask=offs_m[:, None] < L) - + delta_s = tl.sum(o_s * do_s, axis=1).to(DELTAS.type.element_ty) tl.store(DELTAS + offs_m, delta_s, mask=offs_m < L) @triton.jit def _attn_bwd_dq( - Q, K, V, LSE, DELTAS, - DOS, DQ, LUT, + Q, + K, + V, + LSE, + DELTAS, + DOS, + DQ, + LUT, qk_scale: tl.constexpr, topk: tl.constexpr, L: tl.constexpr, @@ -141,28 +153,36 @@ def _attn_bwd_dq( do_s = tl.load(DOS_ptrs, mask=offs_m[:, None] < L) delta_s = tl.load(DELTAS_ptrs, mask=offs_m < L) lse = tl.load(LSE_ptrs, mask=offs_m < L, other=float("inf")) - + dq = tl.zeros([BLOCK_M, D], dtype=tl.float32) for block_idx in tl.range(topk, num_stages=2): idx_n = tl.load(LUT_ptr + block_idx) n_mask = offs_n < L - idx_n * BLOCK_N - + k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None]) v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None]) qk = tl.dot(q, k.T) * (qk_scale * 1.4426950408889634) p = tl.math.exp2(qk - lse[:, None]) p = tl.where(n_mask[None, :], p, 0.0) - + dp = tl.dot(do_s, v.T).to(tl.float32) ds = p * (dp - delta_s[:, None]) dq += tl.dot(ds.to(k.dtype), k) tl.store(DQ_ptrs, dq * qk_scale, mask=offs_m[:, None] < L) - + @triton.jit def _attn_bwd_dkdv( - Q, K, V, DOS, DK, DV, - qk_scale, KBID, LSE, DELTAS, + Q, + K, + V, + DOS, + DK, + DV, + qk_scale, + KBID, + LSE, + DELTAS, L: tl.constexpr, M_BLOCKS: tl.constexpr, N_BLOCKS: tl.constexpr, @@ -196,7 +216,7 @@ def _attn_bwd_dkdv( k = tl.load(K_ptrs, mask=offs_n[:, None] < L) v = tl.load(V_ptrs, mask=offs_n[:, None] < L) - + dk = tl.zeros([BLOCK_N, D], dtype=tl.float32) dv = tl.zeros([BLOCK_N, D], dtype=tl.float32) for idx_m in tl.range(0, L, BLOCK_M2): @@ -215,7 +235,7 @@ def _attn_bwd_dkdv( dpT = tl.dot(v, tl.trans(do)) dsT = pT * (dpT - delta[None, :]) dk += tl.dot(dsT.to(q.dtype), q) - + Q_ptrs += BLOCK_M2 * D DOS_ptrs += BLOCK_M2 * D LSE_ptrs += BLOCK_M2 @@ -225,11 +245,11 @@ def _attn_bwd_dkdv( tl.store(DK_ptrs, dk * qk_scale, mask=offs_n[:, None] < L) tl.store(DV_ptrs, dv, mask=offs_n[:, None] < L) - + class _attention(torch.autograd.Function): """Sparse attention forward/backward with autograd support.""" - + @staticmethod def forward(ctx, q, k, v, k_block_id, lut, topk, BLOCK_M, BLOCK_N, qk_scale=None): assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() @@ -248,15 +268,22 @@ class _attention(torch.autograd.Function): lse = torch.empty(q.shape[:-1], device=q.device, dtype=torch.float32) grid = (M_BLOCKS, B * H) - _attn_fwd[grid]( - q, k, v, qk_scale, topk, - lut, lse, o_s, - L, M_BLOCKS, - D, BLOCK_M, BLOCK_N, - num_warps=4 if q.shape[-1] == 64 else 8, - num_stages=3 - ) - + _attn_fwd[grid](q, + k, + v, + qk_scale, + topk, + lut, + lse, + o_s, + L, + M_BLOCKS, + D, + BLOCK_M, + BLOCK_N, + num_warps=4 if q.shape[-1] == 64 else 8, + num_stages=3) + ctx.save_for_backward(q, k, v, k_block_id, lut, lse, o_s) ctx.qk_scale = qk_scale ctx.topk = topk @@ -282,30 +309,52 @@ class _attention(torch.autograd.Function): grid = (M_BLOCKS, B * H) _attn_bwd_preprocess[grid]( - o_s, do_s, delta_s, - L, D, BLOCK_M, + o_s, + do_s, + delta_s, + L, + D, + BLOCK_M, ) grid = (M_BLOCKS, B * H) - _attn_bwd_dq[grid]( - q, k, v, lse, delta_s, - do_s, dq, lut, - ctx.qk_scale, ctx.topk, - L, M_BLOCKS, - D, BLOCK_M, BLOCK_N, - num_warps=4 if q.shape[-1] == 64 else 8, - num_stages=4 if q.shape[-1] == 64 else 5 - ) + _attn_bwd_dq[grid](q, + k, + v, + lse, + delta_s, + do_s, + dq, + lut, + ctx.qk_scale, + ctx.topk, + L, + M_BLOCKS, + D, + BLOCK_M, + BLOCK_N, + num_warps=4 if q.shape[-1] == 64 else 8, + num_stages=4 if q.shape[-1] == 64 else 5) grid = (N_BLOCKS, B * H) - _attn_bwd_dkdv[grid]( - q, k, v, do_s, dk, dv, - ctx.qk_scale, k_block_id, lse, delta_s, - L, M_BLOCKS, N_BLOCKS, - D, BLOCK_M, BLOCK_N, - BLOCK_SLICE_FACTOR=BLOCK_M // 64, - num_warps=4 if q.shape[-1] == 64 else 8, - num_stages=4 if q.shape[-1] == 64 else 5 - ) + _attn_bwd_dkdv[grid](q, + k, + v, + do_s, + dk, + dv, + ctx.qk_scale, + k_block_id, + lse, + delta_s, + L, + M_BLOCKS, + N_BLOCKS, + D, + BLOCK_M, + BLOCK_N, + BLOCK_SLICE_FACTOR=BLOCK_M // 64, + num_warps=4 if q.shape[-1] == 64 else 8, + num_stages=4 if q.shape[-1] == 64 else 5) return dq, dk, dv, None, None, None, None, None, None diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/st_attn_triton.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/st_attn_triton.py index c086557a1..4ea18f2d2 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/st_attn_triton.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/st_attn_triton.py @@ -3,6 +3,7 @@ import torch import triton import triton.language as tl + def is_cuda(): return triton.runtime.driver.active.get_current_target().backend == "cuda" @@ -14,7 +15,8 @@ def is_hip(): def is_cdna3_cdna4(): target = triton.runtime.driver.active.get_current_target() - return (target.arch == 'gfx1201' or target.arch == 'gfx1101' or target.arch == 'gfx1100' or target.arch == 'gfx1030') + return (target.arch == 'gfx1201' or target.arch == 'gfx1101' or target.arch == 'gfx1100' + or target.arch == 'gfx1030') def get_common_autotune_config(): @@ -56,10 +58,17 @@ def clamp_int(value, min_val, max_val): @triton.jit def _attn_fwd_loop( - q, k, v, kv_mask, m, l, acc, sm_scale, + q, + k, + v, + kv_mask, + m, + l, + acc, + sm_scale, MASK_KV: tl.constexpr, ): - scores = tl.dot(q, k.T) #[BLOCK_Q, BLOCK_KV] + scores = tl.dot(q, k.T) #[BLOCK_Q, BLOCK_KV] scores = scores * sm_scale if MASK_KV: scores = tl.where(kv_mask[None, :], scores, -float('inf')) @@ -86,13 +95,25 @@ def _attn_fwd_loop( ) @triton.jit def triton_sta_kernel( - Q, K, V, output, - batch_size: int, num_heads: int, seq_len: int, head_dim: int, + Q, + K, + V, + output, + batch_size: int, + num_heads: int, + seq_len: int, + head_dim: int, img_seq_len: int, text_length: int, - canvas_t: int, canvas_h: int, canvas_w: int, - kernel_t: int, kernel_h: int, kernel_w: int, - tile_t: int, tile_h: int, tile_w: int, + canvas_t: int, + canvas_h: int, + canvas_w: int, + kernel_t: int, + kernel_h: int, + kernel_w: int, + tile_t: int, + tile_h: int, + tile_w: int, scale: float, has_text: tl.constexpr, text_q: tl.constexpr, @@ -111,8 +132,8 @@ def triton_sta_kernel( q_tile_flat = tl.program_id(2) // q_block_per_tile q_block_idx = tl.program_id(2) % q_block_per_tile - m = tl.full((BLOCK_Q,), -float('inf'), dtype=tl.float32) - l = tl.zeros((BLOCK_Q,), dtype=tl.float32) + m = tl.full((BLOCK_Q, ), -float('inf'), dtype=tl.float32) + l = tl.zeros((BLOCK_Q, ), dtype=tl.float32) acc = tl.zeros((BLOCK_Q, BLOCK_DIM), dtype=tl.float32) q_offset = (batch_idx * num_heads + head_idx) * seq_len * head_dim @@ -125,11 +146,9 @@ def triton_sta_kernel( q_idx = q_base_idx + q_offset_in_tile q_mask = (q_block_idx * BLOCK_Q + tl.arange(0, BLOCK_Q)) < total_tile_size - q = tl.load( - Q + q_offset + q_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], - mask=q_mask[:, None], - other=0.0 - ) # [BLOCK_Q, BLOCK_DIM] + q = tl.load(Q + q_offset + q_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], + mask=q_mask[:, None], + other=0.0) # [BLOCK_Q, BLOCK_DIM] # Scale sm_scale by log_2(e) and use 2^x instead of exp sm_scale = scale * 1.4426950408889634 @@ -175,7 +194,8 @@ def triton_sta_kernel( for kv_tile_t in tl.range(kv_tile_start_t, kv_tile_end_t): for kv_tile_h in tl.range(kv_tile_start_h, kv_tile_end_h): for kv_tile_w in tl.range(kv_tile_start_w, kv_tile_end_w): - kv_base_idx = (kv_tile_t * num_tiles_h * num_tiles_w + kv_tile_h * num_tiles_w + kv_tile_w) * total_tile_size + kv_base_idx = (kv_tile_t * num_tiles_h * num_tiles_w + kv_tile_h * num_tiles_w + + kv_tile_w) * total_tile_size for kv_block_idx in tl.range(0, total_tile_size, BLOCK_KV): kv_offset_in_block = tl.arange(0, BLOCK_KV) @@ -184,20 +204,15 @@ def triton_sta_kernel( kv_offset = (batch_idx * num_heads + head_idx) * seq_len * head_dim - k = tl.load( - K + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], - mask=kv_mask[:, None], - other=0.0 - ) # [BLOCK_KV, BLOCK_DIM] - v = tl.load( - V + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], - mask=kv_mask[:, None], - other=0.0 - ) # [BLOCK_KV, BLOCK_DIM] + k = tl.load(K + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], + mask=kv_mask[:, None], + other=0.0) # [BLOCK_KV, BLOCK_DIM] + v = tl.load(V + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], + mask=kv_mask[:, None], + other=0.0) # [BLOCK_KV, BLOCK_DIM] m, l, acc = _attn_fwd_loop(q, k, v, kv_mask, m, l, acc, sm_scale, False) - # for kv_text if has_text: kv_base_idx = img_seq_len @@ -208,35 +223,32 @@ def triton_sta_kernel( kv_offset = (batch_idx * num_heads + head_idx) * seq_len * head_dim - k = tl.load( - K + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], - mask=kv_mask[:, None], - other=0.0 - ) # [BLOCK_KV, BLOCK_DIM] - v = tl.load( - V + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], - mask=kv_mask[:, None], - other=0.0 - ) # [BLOCK_KV, BLOCK_DIM] + k = tl.load(K + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], + mask=kv_mask[:, None], + other=0.0) # [BLOCK_KV, BLOCK_DIM] + v = tl.load(V + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], + mask=kv_mask[:, None], + other=0.0) # [BLOCK_KV, BLOCK_DIM] m, l, acc = _attn_fwd_loop(q, k, v, kv_mask, m, l, acc, sm_scale, True) - output_acc = acc / l[:, None] - tl.store( - output + q_offset + q_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], - output_acc, - mask=q_mask[:, None] - ) # [BLOCK_Q, BLOCK_DIM] + tl.store(output + q_offset + q_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], + output_acc, + mask=q_mask[:, None]) # [BLOCK_Q, BLOCK_DIM] -def sliding_tile_attention_triton( - q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - window_size, text_length: int, - has_text=True, dit_seq_shape='30x48x80') -> torch.Tensor: +def sliding_tile_attention_triton(q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + window_size, + text_length: int, + has_text=True, + dit_seq_shape='30x48x80') -> torch.Tensor: seq_length = q.shape[2] if has_text: - assert q.shape[2] >= 115200 and q.shape[2] <= 115456, f"Unsupported {dit_seq_shape}, current shape is {q.shape}, only support '30x48x80' for HunyuanVideo" + assert q.shape[2] >= 115200 and q.shape[ + 2] <= 115456, f"Unsupported {dit_seq_shape}, current shape is {q.shape}, only support '30x48x80' for HunyuanVideo" target_size = math.ceil(seq_length / 384) * 384 pad_size = target_size - seq_length if pad_size > 0: @@ -244,22 +256,24 @@ def sliding_tile_attention_triton( k = torch.cat([k, k[:, :, -pad_size:]], dim=2) v = torch.cat([v, v[:, :, -pad_size:]], dim=2) else: - if dit_seq_shape == '36x48x48': # Stepvideo + if dit_seq_shape == '36x48x48': # Stepvideo assert q.shape[2] == 82944 - elif dit_seq_shape == '18x48x80': # Wan + elif dit_seq_shape == '18x48x80': # Wan assert q.shape[2] == 69120 else: - raise ValueError(f"Unsupported {dit_seq_shape}, current shape is {q.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan") + raise ValueError( + f"Unsupported {dit_seq_shape}, current shape is {q.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan" + ) assert q.shape[1] == len(window_size), "Number of heads must match the number of window sizes" batch_size, num_heads, seq_len, head_dim = q.shape - if dit_seq_shape == '30x48x80': # Hunyuan + if dit_seq_shape == '30x48x80': # Hunyuan canvas_t, canvas_h, canvas_w = 30, 48, 80 tile_t, tile_h, tile_w = 6, 8, 8 - elif dit_seq_shape == '36x48x48': # Stepvideo + elif dit_seq_shape == '36x48x48': # Stepvideo canvas_t, canvas_h, canvas_w = 36, 48, 48 tile_t, tile_h, tile_w = 6, 8, 8 - elif dit_seq_shape == '18x48x80': # Wan + elif dit_seq_shape == '18x48x80': # Wan canvas_t, canvas_h, canvas_w = 18, 48, 80 tile_t, tile_h, tile_w = 6, 8, 8 @@ -283,22 +297,36 @@ def sliding_tile_attention_triton( # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs for head_index, (kernel_t, kernel_h, kernel_w) in enumerate(window_size): for batch in range(batch_size): - q_head, k_head, v_head, o_head = (q[batch:batch + 1, head_index:head_index + 1], - k[batch:batch + 1, head_index:head_index + 1], - v[batch:batch + 1, head_index:head_index + 1], - output[batch:batch + 1, head_index:head_index + 1]) + q_head, k_head, v_head, o_head = (q[batch:batch + 1, + head_index:head_index + 1], k[batch:batch + 1, + head_index:head_index + 1], + v[batch:batch + 1, + head_index:head_index + 1], output[batch:batch + 1, + head_index:head_index + 1]) # triton_sta_kernel[(1, 1, num_tiles * triton.cdiv(total_tile_size, BLOCK_Q))]( grid = lambda META: (1, 1, num_tiles * triton.cdiv(total_tile_size, META['BLOCK_Q'])) triton_sta_kernel[grid]( - q_head, k_head, v_head, o_head, - 1, 1, seq_len, head_dim, + q_head, + k_head, + v_head, + o_head, + 1, + 1, + seq_len, + head_dim, img_seq_len, text_length, - canvas_t, canvas_h, canvas_w, - kernel_t, kernel_h, kernel_w, - tile_t, tile_h, tile_w, - scale=1.0 / (head_dim ** 0.5), + canvas_t, + canvas_h, + canvas_w, + kernel_t, + kernel_h, + kernel_w, + tile_t, + tile_h, + tile_w, + scale=1.0 / (head_dim**0.5), has_text=has_text, text_q=False, # BLOCK_Q=BLOCK_Q, @@ -312,15 +340,27 @@ def sliding_tile_attention_triton( # triton_sta_kernel[(batch_size, num_heads, triton.cdiv(total_tile_size, BLOCK_Q))]( grid = lambda META: (batch_size, num_heads, triton.cdiv(total_tile_size, META['BLOCK_Q'])) triton_sta_kernel[grid]( - q, k, v, output, - batch_size, num_heads, seq_len, head_dim, + q, + k, + v, + output, + batch_size, + num_heads, + seq_len, + head_dim, img_seq_len, text_length, - canvas_t, canvas_h, canvas_w, - 3, 3, 3, + canvas_t, + canvas_h, + canvas_w, + 3, + 3, + 3, #kernel_t, kernel_h, kernel_w, - tile_t, tile_h, tile_w, - scale=1.0 / (head_dim ** 0.5), + tile_t, + tile_h, + tile_w, + scale=1.0 / (head_dim**0.5), has_text=has_text, text_q=True, # BLOCK_Q=BLOCK_Q, @@ -332,4 +372,4 @@ def sliding_tile_attention_triton( if pad_size > 0: output = output[:, :, :seq_length] - return output \ No newline at end of file + return output diff --git a/fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py b/fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py index 3d93770ad..092e2f55f 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py @@ -20,6 +20,7 @@ except ImportError: rms_norm_cuda = None layer_norm_cuda = None + def int8_quant(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """ Quantize a floating-point tensor to int8 using a custom CUDA kernel. @@ -67,6 +68,7 @@ def int8_linear( gemm_cuda(x_q, x_s, w_q, w_s, y) return y.reshape(*shape[:-1], n) + def flatten_if_batched(*tensors): """ Flattens all input tensors from (B, N, D_i) to (B * N, D_i) if they are batched (3D). @@ -92,9 +94,8 @@ def flatten_if_batched(*tensors): batched = True batch_size = first.shape[0] assert all(t.shape[0] == batch_size for t in tensors), "All input tensors must have the same batch size" - assert all( - t.shape[1] == first.shape[1] for t in tensors - ), "All input tensors must have the same sequence length" + assert all(t.shape[1] == first.shape[1] + for t in tensors), "All input tensors must have the same sequence length" flat_tensors = [t.reshape(-1, t.shape[-1]) for t in tensors] else: batched = False @@ -167,21 +168,21 @@ def rmsnorm(x, w, eps): # allocate output M, N = x.shape y = torch.empty_like(x, dtype=x.dtype) - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + rstd = torch.empty((M, ), dtype=torch.float32, device=x.device) # heuristics for number of warps num_warps = 8 - + # Avoid illegal memory access N2 = triton.next_power_of_2(N) - + if N <= 512: BLOCK_M = 32 else: BLOCK_M = 1 # Call the triton kernel - _rms_norm_fwd_fused[(triton.cdiv(M, BLOCK_M),)]( # + _rms_norm_fwd_fused[(triton.cdiv(M, BLOCK_M), )]( # x, y, w, @@ -201,6 +202,7 @@ def rmsnorm(x, w, eps): return y, rstd + @triton.jit def _layer_norm_param_fwd_fused( X, # pointer to the input @@ -245,7 +247,7 @@ def _layer_norm_param_fwd_fused( w = tl.load(W + cols) b = tl.load(B + cols) - + x_hat = x_hat * w + b # Write output @@ -260,20 +262,20 @@ def layernorm_param(x, w, b, eps): # allocate output M, N = x.shape y = torch.empty_like(x, dtype=torch.float32) - mean = torch.empty((M,), dtype=torch.float32, device=x.device) - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + mean = torch.empty((M, ), dtype=torch.float32, device=x.device) + rstd = torch.empty((M, ), dtype=torch.float32, device=x.device) # heuristics for number of warps num_warps = 8 N2 = triton.next_power_of_2(N) - + if N <= 512: BLOCK_M = 32 else: BLOCK_M = 1 # enqueue kernel - _layer_norm_param_fwd_fused[(triton.cdiv(M, BLOCK_M),)]( # + _layer_norm_param_fwd_fused[(triton.cdiv(M, BLOCK_M), )]( # x, y, w, @@ -355,20 +357,20 @@ def layernorm_noparam(x, eps): # allocate output M, N = x.shape y = torch.empty_like(x, dtype=torch.float32) - mean = torch.empty((M,), dtype=torch.float32, device=x.device) - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + mean = torch.empty((M, ), dtype=torch.float32, device=x.device) + rstd = torch.empty((M, ), dtype=torch.float32, device=x.device) # heuristics for number of warps num_warps = 8 N2 = triton.next_power_of_2(N) - + if N <= 512: BLOCK_M = 32 else: BLOCK_M = 1 # enqueue kernel - _layer_norm_noparam_fwd_fused[(triton.cdiv(M, BLOCK_M),)]( # + _layer_norm_noparam_fwd_fused[(triton.cdiv(M, BLOCK_M), )]( # x, y, mean, @@ -388,6 +390,7 @@ def layernorm_noparam(x, eps): return y, mean, rstd + def layernorm(x, w, b, eps, elementwise_affine=True): if elementwise_affine: assert w is not None and b is not None @@ -396,9 +399,11 @@ def layernorm(x, w, b, eps, elementwise_affine=True): assert w is None and b is None return layernorm_noparam(x, eps) + def cdiv(a: int, b: int): return (a + b - 1) // b + def dequantize_weight(w_q, w_s, dtype): # w_q: (N, K) int8 # w_s: (NB, KB) float32 @@ -412,12 +417,14 @@ def dequantize_weight(w_q, w_s, dtype): w_s_exp = w_s_exp[:N, :K] return w_q.to(dtype) * w_s_exp.to(dtype) + class Int8LinearFunction(torch.autograd.Function): + @staticmethod def forward(ctx, x, w_q, w_s, bias=None): ctx.save_for_backward(x, w_q, w_s, bias) ctx.bias_requires_grad = bias.requires_grad if bias is not None else False - + out = int8_linear(x, w_q, w_s) if bias is not None: out = out + bias @@ -426,24 +433,26 @@ class Int8LinearFunction(torch.autograd.Function): @staticmethod def backward(ctx, grad_output): x, w_q, w_s, bias = ctx.saved_tensors - + grad_input = None - grad_weight = None + grad_weight = None grad_scale = None grad_bias = None - - if ctx.needs_input_grad[0]: + + if ctx.needs_input_grad[0]: # grad_input = grad_output @ W w_float = dequantize_weight(w_q, w_s, grad_output.dtype) grad_input = torch.matmul(grad_output, w_float) - + if ctx.bias_requires_grad: dim_to_sum = list(range(grad_output.dim() - 1)) grad_bias = grad_output.sum(dim=dim_to_sum) - + return grad_input, grad_weight, grad_scale, grad_bias + class RMSNormFunction(torch.autograd.Function): + @staticmethod def forward(ctx, x, w, eps): y, rstd = rmsnorm(x, w, eps) @@ -459,14 +468,14 @@ class RMSNormFunction(torch.autograd.Function): x, w, rstd = ctx.saved_tensors eps = ctx.eps N = x.shape[-1] - + # dL/dy * w dx = grad_output * w - + # Expand rstd - rstd = rstd.unsqueeze(-1) # (B, 1) or (M, 1) + rstd = rstd.unsqueeze(-1) # (B, 1) or (M, 1) if x.dim() == 3: - # Flatten if necessary or handle dimensions. + # Flatten if necessary or handle dimensions. # forward flattens x to (M, N) internally but saved x might be 3D? # rmsnorm takes x, flattens it. But does it modify x in place? No. # But ctx.save_for_backward saves the original x (3D if input was 3D). @@ -479,36 +488,36 @@ class RMSNormFunction(torch.autograd.Function): x_flat = x grad_output_flat = grad_output dx = dx - + # Standard RMSNorm backward # dy = grad_output_flat # x_hat = x * rstd # w * dy # c1 = mean(dy * w * x) * rstd^2 # dx = (dy * w - c1 * x) * rstd - + # More precise: # y = x * rstd * w # dy = grad_output # dw = sum(dy * x * rstd, dim=0) - + grad_w = None - if ctx.needs_input_grad[1]: # w - # grad_w = (grad_output_flat * x_flat * rstd).sum(dim=0) - # More accurate gradient for weight when considering rstd was computed from x - # Actually, for Affine part (w), it is just sum(dL/dy * x_hat). - # y = x_hat * w - # dL/dw = sum(dL/dy * x_hat) - x_hat = x_flat * rstd - grad_w = (grad_output_flat * x_hat).sum(dim=0) + if ctx.needs_input_grad[1]: # w + # grad_w = (grad_output_flat * x_flat * rstd).sum(dim=0) + # More accurate gradient for weight when considering rstd was computed from x + # Actually, for Affine part (w), it is just sum(dL/dy * x_hat). + # y = x_hat * w + # dL/dw = sum(dL/dy * x_hat) + x_hat = x_flat * rstd + grad_w = (grad_output_flat * x_hat).sum(dim=0) grad_input = None - if ctx.needs_input_grad[0]: # x + if ctx.needs_input_grad[0]: # x # x_hat = x * rstd - # dx = rstd * (w * dy - mean(w * dy * x_hat) * x_hat) + # dx = rstd * (w * dy - mean(w * dy * x_hat) * x_hat) # but for RMSNorm: # dx = rstd * (w * dy - (x * rstd^2) * mean(w * dy * x)) -> check formula - + # Using PyTorch autograd for reference logic: # y = x / sqrt(mean(x^2) + eps) * w # Let sigma = sqrt(...) @@ -518,12 +527,12 @@ class RMSNormFunction(torch.autograd.Function): # dL/dx = (dL/dy * w)/sigma - sum(dL/dy * w * x) * x / (N * sigma^3) # = (1/sigma) * [ (dL/dy * w) - x * sum(dL/dy * w * x) / (N * sigma^2) ] # = rstd * [ (dL/dy * w) - x * rstd^2 * mean(dL/dy * w * x) ] # Wait, sum/N is mean - + # Let's compute: dy_w = grad_output_flat * w # term2 = (dy_w * x_flat).sum(dim=-1, keepdim=True) / N # mean(dy*w*x) # grad_input = rstd * (dy_w - x_flat * (rstd ** 2) * term2) - + # Re-derivation: # x_hat = x * rstd # y = x_hat * w @@ -531,20 +540,22 @@ class RMSNormFunction(torch.autograd.Function): # dy/dx = w * dx_hat/dx # dx_hat/dx = rstd * (I - x * x^T * rstd^2 / N) ?? No. # dx_hat_i / dx_j = rstd * (delta_ij - x_i * x_j * rstd^2 / N) - + # dL/dx_hat = dL/dy * w # dL/dx = rstd * (dL/dx_hat - x * rstd^2 * mean(dL/dx_hat * x)) - + dx_hat = grad_output_flat * w dx_hat_x_mean = (dx_hat * x_flat).mean(dim=-1, keepdim=True) - grad_input = rstd * (dx_hat - x_flat * (rstd ** 2) * dx_hat_x_mean) - + grad_input = rstd * (dx_hat - x_flat * (rstd**2) * dx_hat_x_mean) + if x.dim() == 3: grad_input = grad_input.reshape(x.shape) - + return grad_input, grad_w, None + class LayerNormFunction(torch.autograd.Function): + @staticmethod def forward(ctx, x, w, b, eps, elementwise_affine): if elementwise_affine: @@ -553,7 +564,7 @@ class LayerNormFunction(torch.autograd.Function): else: y, mean, rstd = layernorm_noparam(x, eps) ctx.save_for_backward(x, None, None, mean, rstd) - + ctx.eps = eps ctx.elementwise_affine = elementwise_affine return y @@ -562,60 +573,62 @@ class LayerNormFunction(torch.autograd.Function): def backward(ctx, grad_output): x, w, b, mean, rstd = ctx.saved_tensors N = x.shape[-1] - + if x.dim() == 3: x_flat = x.reshape(-1, N) grad_output_flat = grad_output.reshape(-1, N) else: x_flat = x grad_output_flat = grad_output - + if w is None: - w_eff = 1.0 + w_eff = 1.0 else: - w_eff = w - + w_eff = w + # dx calculation # x_hat = (x - mean) * rstd # y = x_hat * w + b # dL/dx_hat = dL/dy * w - + dy = grad_output_flat dx_hat = dy * w_eff - + # dL/dvar = sum(dL/dx_hat * (x-mean) * (-0.5) * (var+eps)^(-1.5)) # = -0.5 * rstd^3 * sum(dx_hat * (x-mean)) # dL/dmean = sum(dL/dx_hat * (-rstd)) + dL/dvar * (-2/N) * sum(x-mean) # term sum(x-mean) is 0. So second part vanishes. # dL/dmean = -rstd * sum(dx_hat) - + # dL/dx = dL/dx_hat * rstd + dL/dvar * 2(x-mean)/N + dL/dmean * 1/N # = dx_hat * rstd + (-0.5 * rstd^3 * sum(dx_hat * (x-mean))) * 2(x-mean)/N + (-rstd * sum(dx_hat))/N # = rstd * [ dx_hat - mean(dx_hat) - (x-mean)*rstd^2 * mean(dx_hat * (x-mean)) ] - + x_centered = x_flat - mean.unsqueeze(-1) dx_hat_mean = dx_hat.mean(dim=-1, keepdim=True) dx_hat_x_centered_mean = (dx_hat * x_centered).mean(dim=-1, keepdim=True) - - grad_input = rstd.unsqueeze(-1) * (dx_hat - dx_hat_mean - x_centered * (rstd.unsqueeze(-1)**2) * dx_hat_x_centered_mean) - + + grad_input = rstd.unsqueeze(-1) * (dx_hat - dx_hat_mean - x_centered * + (rstd.unsqueeze(-1)**2) * dx_hat_x_centered_mean) + if x.dim() == 3: grad_input = grad_input.reshape(x.shape) grad_w = None grad_b = None - + if ctx.elementwise_affine: if ctx.needs_input_grad[1]: - x_hat = x_centered * rstd.unsqueeze(-1) - grad_w = (dy * x_hat).sum(dim=0) + x_hat = x_centered * rstd.unsqueeze(-1) + grad_w = (dy * x_hat).sum(dim=0) if ctx.needs_input_grad[2]: - grad_b = dy.sum(dim=0) - + grad_b = dy.sum(dim=0) + return grad_input, grad_w, grad_b, None, None class Int8Linear(nn.Module): + def __init__(self, in_features, out_features, bias=True, dtype=torch.bfloat16): super().__init__() self.in_features = in_features @@ -623,27 +636,24 @@ class Int8Linear(nn.Module): row_blocks = cdiv(out_features, b=128) col_blocks = cdiv(in_features, b=128) - + self.register_buffer("int8_weight", torch.empty((out_features, in_features), dtype=torch.int8)) self.register_buffer("scale", torch.empty((row_blocks, col_blocks), dtype=torch.float32)) if bias: self.register_buffer("bias", torch.empty(out_features, dtype=dtype)) else: self.bias = None - def forward(self, x): return Int8LinearFunction.apply(x, self.int8_weight, self.scale, self.bias) @classmethod def from_linear(cls, original_linear: nn.Linear, quantize: bool = True): - - int8_layer = cls( - original_linear.in_features, - original_linear.out_features, - bias=original_linear.bias is not None, - dtype=original_linear.weight.dtype - ) + + int8_layer = cls(original_linear.in_features, + original_linear.out_features, + bias=original_linear.bias is not None, + dtype=original_linear.weight.dtype) if quantize: w_data = original_linear.weight.data.cuda() int8_w, scale = int8_quant(w_data) @@ -652,10 +662,12 @@ class Int8Linear(nn.Module): int8_layer.scale.copy_(scale) if original_linear.bias is not None: int8_layer.bias.data.copy_(original_linear.bias.data.cuda()) - + return int8_layer - + + class FastRMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-5): super().__init__() self.dim = dim @@ -667,23 +679,15 @@ class FastRMSNorm(nn.Module): @classmethod def from_rmsnorm(cls, original_rmsnorm): - rmsnorm_layer = cls( - dim=original_rmsnorm.dim, - eps=original_rmsnorm.eps - ) + rmsnorm_layer = cls(dim=original_rmsnorm.dim, eps=original_rmsnorm.eps) if original_rmsnorm.weight.device != torch.device('meta'): rmsnorm_layer.weight.data.copy_(original_rmsnorm.weight.float().data) return rmsnorm_layer - + + class FastLayerNorm(nn.Module): - def __init__( - self, - dim: int, - eps: float = 1e-5, - elementwise_affine: bool = False, - bias: bool = True - ) : + def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine: bool = False, bias: bool = True): super().__init__() self.dim = dim # type: ignore[arg-type] self.eps = eps @@ -700,15 +704,13 @@ class FastLayerNorm(nn.Module): def forward(self, x): return LayerNormFunction.apply(x.float(), self.weight, self.bias, self.eps, self.elementwise_affine).to(x.dtype) - + @classmethod def from_layernorm(cls, original_layernorm): - layernorm_layer = cls( - dim=original_layernorm.normalized_shape[0], - eps=original_layernorm.eps, - elementwise_affine=False if original_layernorm.weight is None else True, - bias=original_layernorm.bias is not None - ) + layernorm_layer = cls(dim=original_layernorm.normalized_shape[0], + eps=original_layernorm.eps, + elementwise_affine=False if original_layernorm.weight is None else True, + bias=original_layernorm.bias is not None) if original_layernorm.weight is not None and original_layernorm.weight.device != torch.device('meta'): layernorm_layer.weight.data.copy_(original_layernorm.weight.data) if original_layernorm.bias is not None and original_layernorm.bias.device != torch.device('meta'): diff --git a/fastvideo-kernel/python/fastvideo_kernel/vmoba.py b/fastvideo-kernel/python/fastvideo_kernel/vmoba.py index 7eb1c64f3..4149a59f0 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/vmoba.py +++ b/fastvideo-kernel/python/fastvideo_kernel/vmoba.py @@ -10,9 +10,7 @@ try: except ImportError: def _unsupported(*args, **kwargs): - raise ImportError( - "flash-attn is not installed. Please install it, e.g., `uv pip install flash-attn`." - ) + raise ImportError("flash-attn is not installed. Please install it, e.g., `uv pip install flash-attn`.") _flash_attn_varlen_forward = _unsupported _flash_attn_varlen_backward = _unsupported @@ -39,25 +37,17 @@ def calc_chunks(cu_seqlen, moba_chunk_size): ) cu_num_chunk[1:] = batch_num_chunk.cumsum(dim=0) num_chunk = cu_num_chunk[-1] - chunk_sizes = torch.full((num_chunk + 1, ), - moba_chunk_size, - dtype=torch.int32, - device=cu_seqlen.device) + chunk_sizes = torch.full((num_chunk + 1, ), moba_chunk_size, dtype=torch.int32, device=cu_seqlen.device) chunk_sizes[0] = 0 - batch_last_chunk_size = batch_sizes - (batch_num_chunk - - 1) * moba_chunk_size + batch_last_chunk_size = batch_sizes - (batch_num_chunk - 1) * moba_chunk_size chunk_sizes[cu_num_chunk[1:]] = batch_last_chunk_size cu_chunk = chunk_sizes.cumsum(dim=-1, dtype=torch.int32) - chunk_to_batch = torch.zeros((num_chunk, ), - dtype=torch.int32, - device=cu_seqlen.device) + chunk_to_batch = torch.zeros((num_chunk, ), dtype=torch.int32, device=cu_seqlen.device) chunk_to_batch[cu_num_chunk[1:-1]] = 1 chunk_to_batch = chunk_to_batch.cumsum(dim=0, dtype=torch.int32) # Do not filter out any chunk - filtered_chunk_indices = torch.arange(num_chunk, - device=cu_seqlen.device, - dtype=torch.int32) + filtered_chunk_indices = torch.arange(num_chunk, device=cu_seqlen.device, dtype=torch.int32) num_filtered_chunk = num_chunk return cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch @@ -66,9 +56,7 @@ def calc_chunks(cu_seqlen, moba_chunk_size): # --- Threshold Selection Helper Functions --- -def _select_threshold_query_head(gate: torch.Tensor, - valid_gate_mask: torch.Tensor, - gate_self_chunk_mask: torch.Tensor, +def _select_threshold_query_head(gate: torch.Tensor, valid_gate_mask: torch.Tensor, gate_self_chunk_mask: torch.Tensor, simsum_threshold: float) -> torch.Tensor: """ Selects chunks for each pair based on threshold. @@ -78,16 +66,13 @@ def _select_threshold_query_head(gate: torch.Tensor, eps = 1e-6 # LSE‐style normalization per (across chunks) - gate_masked = torch.where(valid_gate_mask, gate, - -torch.inf) # Use -inf for max - gate_min_val = torch.where(valid_gate_mask, gate, - torch.inf) # Use +inf for min + gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max + gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min row_min = gate_min_val.amin(dim=0) # (H, S) row_max = gate_masked.amax(dim=0) # (H, S) denom = row_max - row_min - denom = torch.where(denom <= eps, torch.ones_like(denom), - denom) # avoid divide‑by‑zero + denom = torch.where(denom <= eps, torch.ones_like(denom), denom) # avoid divide‑by‑zero gate_norm = (gate - row_min.unsqueeze(0)) / denom.unsqueeze(0) gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S) @@ -97,18 +82,15 @@ def _select_threshold_query_head(gate: torch.Tensor, # 2) compute how much more normalized weight we need beyond self total_norm_sum = gate_norm.sum(dim=0) # (H, S) - remain_ratio = simsum_threshold - self_norm / (total_norm_sum + eps - ) # (H, S) - remain_ratio = torch.clamp(remain_ratio, - min=0.0) # if already ≥ thresh, no extra needed + remain_ratio = simsum_threshold - self_norm / (total_norm_sum + eps) # (H, S) + remain_ratio = torch.clamp(remain_ratio, min=0.0) # if already ≥ thresh, no extra needed # 3) zero out the self‐chunk in a copy, so we only sort “others” others_norm = gate_norm.clone() others_norm[gate_self_chunk_mask] = 0.0 # 4) sort the other chunks by descending norm, per - sorted_norm, sorted_idx = torch.sort(others_norm, descending=True, - dim=0) # (C, H, S) + sorted_norm, sorted_idx = torch.sort(others_norm, descending=True, dim=0) # (C, H, S) # 5) cumulative‑sum the sorted norms per cumsum_others = sorted_norm.cumsum(dim=0) # (C, H, S) @@ -118,9 +100,7 @@ def _select_threshold_query_head(gate: torch.Tensor, cond = ratio >= remain_ratio.unsqueeze(0) # (C, H, S) boolean mask any_cond = cond.any(dim=0) # (H, S) # Find the index of the first True value along dim 0. If none, use C-1. - cutoff = torch.where(any_cond, - cond.float().argmax(dim=0), - torch.full_like(any_cond, fill_value=C - 1)) # (H, S) + cutoff = torch.where(any_cond, cond.float().argmax(dim=0), torch.full_like(any_cond, fill_value=C - 1)) # (H, S) # 7) build a mask in sorted order up to that cutoff idx_range = torch.arange(C, device=gate.device).view(-1, 1, 1) # (C, 1, 1) @@ -136,8 +116,7 @@ def _select_threshold_query_head(gate: torch.Tensor, return final_gate_mask -def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor, - gate_self_chunk_mask: torch.Tensor, +def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor, gate_self_chunk_mask: torch.Tensor, simsum_threshold: float) -> torch.Tensor: """ Selects pairs for each block based on threshold. @@ -148,16 +127,13 @@ def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor, eps = 1e-6 # LSE‐style normalization per block (across heads and queries) - gate_masked = torch.where(valid_gate_mask, gate, - -torch.inf) # Use -inf for max - gate_min_val = torch.where(valid_gate_mask, gate, - torch.inf) # Use +inf for min + gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max + gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min block_max = gate_masked.amax(dim=(1, 2), keepdim=True) # (C, 1, 1) block_min = gate_min_val.amin(dim=(1, 2), keepdim=True) # (C, 1, 1) block_denom = block_max - block_min - block_denom = torch.where(block_denom <= eps, torch.ones_like(block_denom), - block_denom) # (C, 1, 1) + block_denom = torch.where(block_denom <= eps, torch.ones_like(block_denom), block_denom) # (C, 1, 1) gate_norm = (gate - block_min) / block_denom # (C, H, S) gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S) @@ -169,8 +145,7 @@ def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor, # 2) compute how much more normalized weight each block needs beyond its self-chunk contributions total_norm_sum_per_block = gate_norm.sum(dim=(1, 2)) # (C,) - remain_ratio = simsum_threshold - self_norm_sum_per_block / ( - total_norm_sum_per_block + eps) # (C,) + remain_ratio = simsum_threshold - self_norm_sum_per_block / (total_norm_sum_per_block + eps) # (C,) remain_ratio = torch.clamp(remain_ratio, min=0.0) # (C,) # 3) zero out the self‐chunk entries in a copy, so we only sort “others” @@ -179,31 +154,25 @@ def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor, # 4) sort the other pairs by descending norm, per block others_flat = others_norm.contiguous().view(C, HS) # (C, H*S) - sorted_others_flat, sorted_indices_flat = torch.sort( - others_flat, dim=1, descending=True) # (C, H*S) + sorted_others_flat, sorted_indices_flat = torch.sort(others_flat, dim=1, descending=True) # (C, H*S) # 5) cumulative‑sum the sorted norms per block cumsum_others_flat = sorted_others_flat.cumsum(dim=1) # (C, H*S) # 6) for each block, find the smallest k where cumsum_ratio ≥ remain_ratio - ratio_flat = cumsum_others_flat / ( - total_norm_sum_per_block.unsqueeze(1) + eps) # (C, H*S) + ratio_flat = cumsum_others_flat / (total_norm_sum_per_block.unsqueeze(1) + eps) # (C, H*S) cond_flat = ratio_flat >= remain_ratio.unsqueeze(1) # (C, H*S) boolean mask any_cond = cond_flat.any(dim=1) # (C,) # Find the index of the first True value along dim 1. If none, use HS-1. cutoff_flat = torch.where(any_cond, - cond_flat.float().argmax(dim=1), - torch.full_like(any_cond, - fill_value=HS - 1)) # (C,) + cond_flat.float().argmax(dim=1), torch.full_like(any_cond, fill_value=HS - 1)) # (C,) # 7) build a mask in sorted order up to that cutoff per block - idx_range_flat = torch.arange(HS, - device=gate.device).unsqueeze(0) # (1, H*S) + idx_range_flat = torch.arange(HS, device=gate.device).unsqueeze(0) # (1, H*S) sorted_mask_flat = idx_range_flat <= cutoff_flat.unsqueeze(1) # (C, H*S) # 8) scatter it back to original order per block - others_mask_flat = torch.zeros_like(others_flat, - dtype=torch.bool) # (C, H*S) + others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C, H*S) others_mask_flat.scatter_(1, sorted_indices_flat, sorted_mask_flat) others_mask = others_mask_flat.view(C, H, S) # (C, H, S) @@ -213,8 +182,7 @@ def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor, return final_gate_mask -def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor, - gate_self_chunk_mask: torch.Tensor, +def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor, gate_self_chunk_mask: torch.Tensor, simsum_threshold: float) -> torch.Tensor: """ Selects triplets globally based on threshold. @@ -225,17 +193,14 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor, eps = 1e-6 # LSE‐style normalization globally across all valid entries - gate_masked = torch.where(valid_gate_mask, gate, - -torch.inf) # Use -inf for max - gate_min_val = torch.where(valid_gate_mask, gate, - torch.inf) # Use +inf for min + gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max + gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min overall_max = gate_masked.max() # scalar overall_min = gate_min_val.min() # scalar overall_denom = overall_max - overall_min - overall_denom = torch.where( - overall_denom <= eps, - torch.tensor(1.0, device=gate.device, dtype=gate.dtype), overall_denom) + overall_denom = torch.where(overall_denom <= eps, torch.tensor(1.0, device=gate.device, dtype=gate.dtype), + overall_denom) gate_norm = (gate - overall_min) / overall_denom # (C, H, S) gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S) @@ -247,8 +212,7 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor, # 2) compute how much more normalized weight is needed globally beyond self-chunk contributions total_norm_sum_overall = gate_norm.sum() # scalar - remain_ratio = simsum_threshold - self_norm_sum_overall / ( - total_norm_sum_overall + eps) # scalar + remain_ratio = simsum_threshold - self_norm_sum_overall / (total_norm_sum_overall + eps) # scalar remain_ratio = torch.clamp(remain_ratio, min=0.0) # scalar # 3) zero out the self‐chunk entries in a copy, so we only sort “others” @@ -264,18 +228,14 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor, valid_others_indices = torch.where(valid_others_mask_flat)[0] valid_others_values = others_flat[valid_others_indices] - sorted_others_values, sort_perm = torch.sort( - valid_others_values, descending=True) # (N_valid_others,) - sorted_original_indices = valid_others_indices[ - sort_perm] # Original indices in C*H*S space, sorted by value + sorted_others_values, sort_perm = torch.sort(valid_others_values, descending=True) # (N_valid_others,) + sorted_original_indices = valid_others_indices[sort_perm] # Original indices in C*H*S space, sorted by value # 5) cumulative‑sum the sorted valid 'other' norms globally - cumsum_others_values = sorted_others_values.cumsum( - dim=0) # (N_valid_others,) + cumsum_others_values = sorted_others_values.cumsum(dim=0) # (N_valid_others,) # 6) find the smallest k where cumsum_ratio ≥ remain_ratio globally - ratio_values = cumsum_others_values / (total_norm_sum_overall + eps - ) # (N_valid_others,) + ratio_values = cumsum_others_values / (total_norm_sum_overall + eps) # (N_valid_others,) cond_values = ratio_values >= remain_ratio # (N_valid_others,) boolean mask any_cond = cond_values.any() # scalar @@ -283,19 +243,15 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor, cutoff_idx_in_sorted = torch.where( any_cond, cond_values.float().argmax(dim=0), - torch.tensor(len(sorted_others_values) - 1, - device=gate.device, - dtype=torch.long)) + torch.tensor(len(sorted_others_values) - 1, device=gate.device, dtype=torch.long)) # 7) build a mask selecting the top-k others based on the cutoff # Select the original indices corresponding to the top entries in the sorted list selected_other_indices = sorted_original_indices[:cutoff_idx_in_sorted + 1] # 8) create the mask in the original flat shape - others_mask_flat = torch.zeros_like(others_flat, - dtype=torch.bool) # (C*H*S,) - if selected_other_indices.numel( - ) > 0: # Check if any 'other' indices were selected + others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C*H*S,) + if selected_other_indices.numel() > 0: # Check if any 'other' indices were selected others_mask_flat[selected_other_indices] = True others_mask = others_mask_flat.view(C, H, S) # (C, H, S) @@ -305,9 +261,7 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor, return final_gate_mask -def _select_threshold_head_global(gate: torch.Tensor, - valid_gate_mask: torch.Tensor, - gate_self_chunk_mask: torch.Tensor, +def _select_threshold_head_global(gate: torch.Tensor, valid_gate_mask: torch.Tensor, gate_self_chunk_mask: torch.Tensor, simsum_threshold: float) -> torch.Tensor: """ Selects globally for each head based on threshold. @@ -334,8 +288,7 @@ def _select_threshold_head_global(gate: torch.Tensor, total_norm_sum = gate_norm.sum(dim=(0, 2)) # (H,) # 4) how much more normalized weight needed per head - remain_ratio = simsum_threshold - self_norm_sum / (total_norm_sum + eps - ) # (H,) + remain_ratio = simsum_threshold - self_norm_sum / (total_norm_sum + eps) # (H,) remain_ratio = torch.clamp(remain_ratio, min=0.0) # 5) zero out self‐chunk entries to focus on "others" @@ -349,10 +302,8 @@ def _select_threshold_head_global(gate: torch.Tensor, .permute(1, 0, 2).reshape(H, CS) # (H, C*S) # 7) vectorized selection of “others” per head - masked_flat = torch.where(valid_flat, others_flat, - torch.zeros_like(others_flat)) - sorted_vals, sorted_idx = torch.sort(masked_flat, dim=1, - descending=True) # (H, C*S) + masked_flat = torch.where(valid_flat, others_flat, torch.zeros_like(others_flat)) + sorted_vals, sorted_idx = torch.sort(masked_flat, dim=1, descending=True) # (H, C*S) cumsum_vals = sorted_vals.cumsum(dim=1) # (H, C*S) ratio_vals = cumsum_vals / (total_norm_sum.unsqueeze(1) + eps) # (H, C*S) @@ -360,8 +311,7 @@ def _select_threshold_head_global(gate: torch.Tensor, has_cutoff = cond.any(dim=1) # (H,) default = torch.full((H, ), CS - 1, device=gate.device, dtype=torch.long) - cutoff = torch.where(has_cutoff, - cond.float().argmax(dim=1), default) # (H,) + cutoff = torch.where(has_cutoff, cond.float().argmax(dim=1), default) # (H,) idx_range = torch.arange(CS, device=gate.device).unsqueeze(0) # (1, C*S) sorted_mask = idx_range <= cutoff.unsqueeze(1) # (H, C*S) @@ -430,25 +380,19 @@ class MixedAttention(torch.autograd.Function): self_attn_lse_sh = self_attn_lse_hs.t().contiguous() moba_attn_lse = moba_attn_lse_hs.t().contiguous() - output = torch.zeros((q.shape[0], q.shape[1], q.shape[2]), - device=q.device, - dtype=torch.float32) + output = torch.zeros((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) output_2d = output.view(-1, q.shape[2]) max_lse_1d = self_attn_lse_sh.view(-1) - max_lse_1d = max_lse_1d.index_reduce(0, moba_q_sh_indices, - moba_attn_lse.view(-1), "amax") - self_attn_lse_sh = self_attn_lse_sh - max_lse_1d.view_as( - self_attn_lse_sh) - moba_attn_lse = (moba_attn_lse.view(-1).sub( - max_lse_1d.index_select( - 0, moba_q_sh_indices)).reshape_as(moba_attn_lse)) + max_lse_1d = max_lse_1d.index_reduce(0, moba_q_sh_indices, moba_attn_lse.view(-1), "amax") + self_attn_lse_sh = self_attn_lse_sh - max_lse_1d.view_as(self_attn_lse_sh) + moba_attn_lse = (moba_attn_lse.view(-1).sub(max_lse_1d.index_select( + 0, moba_q_sh_indices)).reshape_as(moba_attn_lse)) mixed_attn_se_sh = self_attn_lse_sh.exp() moba_attn_se = moba_attn_lse.exp() - mixed_attn_se_sh.view(-1).index_add_(0, moba_q_sh_indices, - moba_attn_se.view(-1)) + mixed_attn_se_sh.view(-1).index_add_(0, moba_q_sh_indices, moba_attn_se.view(-1)) mixed_attn_lse_sh = mixed_attn_se_sh.log() # Combine self-attention output @@ -457,15 +401,13 @@ class MixedAttention(torch.autograd.Function): output_2d += self_attn_out_sh.reshape_as(output_2d) # Combine MOBA attention output - mixed_attn_lse = (mixed_attn_lse_sh.view(-1).index_select( - 0, moba_q_sh_indices).view_as(moba_attn_lse)) + mixed_attn_lse = (mixed_attn_lse_sh.view(-1).index_select(0, moba_q_sh_indices).view_as(moba_attn_lse)) factor = (moba_attn_lse - mixed_attn_lse).exp() # [S, H] moba_attn_out = moba_attn_out * factor.unsqueeze(-1) raw_attn_out = moba_attn_out.view(-1, moba_attn_out.shape[-1]) output_2d.index_add_(0, moba_q_sh_indices, raw_attn_out) output = output.to(q.dtype) - mixed_attn_lse_sh = mixed_attn_lse_sh + max_lse_1d.view_as( - mixed_attn_se_sh) + mixed_attn_lse_sh = mixed_attn_lse_sh + max_lse_1d.view_as(mixed_attn_se_sh) ctx.save_for_backward( output, mixed_attn_lse_sh, @@ -508,37 +450,33 @@ class MixedAttention(torch.autograd.Function): dq = torch.empty_like(q) dk = torch.empty_like(k) dv = torch.empty_like(v) - _ = _flash_attn_varlen_backward( - dout=d_output, - q=q, - k=k, - v=v, - out=output, - softmax_lse=mixed_attn_vlse_sh.t().contiguous(), - dq=dq, - dk=dk, - dv=dv, - cu_seqlens_q=self_attn_cu_seqlen, - cu_seqlens_k=self_attn_cu_seqlen, - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, - softmax_scale=softmax_scale, - causal=False, - dropout_p=0.0, - softcap=0.0, - alibi_slopes=None, - deterministic=True, - window_size_left=-1, - window_size_right=-1) + _ = _flash_attn_varlen_backward(dout=d_output, + q=q, + k=k, + v=v, + out=output, + softmax_lse=mixed_attn_vlse_sh.t().contiguous(), + dq=dq, + dk=dk, + dv=dv, + cu_seqlens_q=self_attn_cu_seqlen, + cu_seqlens_k=self_attn_cu_seqlen, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + softmax_scale=softmax_scale, + causal=False, + dropout_p=0.0, + softcap=0.0, + alibi_slopes=None, + deterministic=True, + window_size_left=-1, + window_size_right=-1) headdim = q.shape[-1] - d_moba_output = (d_output.view(-1, headdim).index_select( - 0, moba_q_sh_indices).unsqueeze(1)) - moba_output = (output.view(-1, headdim).index_select( - 0, moba_q_sh_indices).unsqueeze(1)) + d_moba_output = (d_output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1)) + moba_output = (output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1)) - mixed_attn_vlse = (mixed_attn_vlse_sh.view(-1).index_select( - 0, moba_q_sh_indices).view(1, -1)) + mixed_attn_vlse = (mixed_attn_vlse_sh.view(-1).index_select(0, moba_q_sh_indices).view(1, -1)) dmq = torch.empty_like(moba_q) dmkv = torch.empty_like(moba_kv) @@ -595,8 +533,7 @@ def moba_attn_varlen( seqlen, num_head, head_dim = q.shape # Compute chunk boundaries. - cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks( - cu_seqlens, moba_chunk_size) + cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks(cu_seqlens, moba_chunk_size) self_attn_cu_seqlen = cu_chunk @@ -609,20 +546,15 @@ def moba_attn_varlen( chunk_lengths = chunk_ends - chunk_starts # [num_filtered_chunk] max_chunk_len = int(chunk_lengths.max().item()) - range_tensor = torch.arange(max_chunk_len, - device=kv.device, - dtype=chunk_starts.dtype).unsqueeze(0) + range_tensor = torch.arange(max_chunk_len, device=kv.device, dtype=chunk_starts.dtype).unsqueeze(0) indices = chunk_starts.unsqueeze(1) + range_tensor indices = torch.clamp(indices, max=kv.shape[0] - 1) valid_mask = range_tensor < chunk_lengths.unsqueeze(1) - gathered = kv[indices.view(-1)].view(num_filtered_chunk, max_chunk_len, - *kv.shape[1:]) - gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze( - -1).type_as(gathered) + gathered = kv[indices.view(-1)].view(num_filtered_chunk, max_chunk_len, *kv.shape[1:]) + gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).type_as(gathered) # Compute key_gate_weight over valid tokens. - key_values = gathered[:, :, 0].float( - ) # [num_filtered_chunk, max_chunk_len, num_head, head_dim] + key_values = gathered[:, :, 0].float() # [num_filtered_chunk, max_chunk_len, num_head, head_dim] valid_mask_exp = valid_mask.unsqueeze(-1).unsqueeze(-1) key_sum = (key_values * valid_mask_exp).sum(dim=1) divisor = valid_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1) @@ -631,19 +563,15 @@ def moba_attn_varlen( # Compute gate logits between key_gate_weight and queries. q_float = q.float() # gate = torch.einsum("nhd,shd->nhs", key_gate_weight, q_float) # [num_filtered_chunk, num_head, seqlen] - gate = torch.bmm(key_gate_weight.permute(1, 0, 2), - q_float.permute(1, 0, 2).transpose(1, 2)).permute(1, 0, 2) + gate = torch.bmm(key_gate_weight.permute(1, 0, 2), q_float.permute(1, 0, 2).transpose(1, 2)).permute(1, 0, 2) # Amplify the diagonal (self chunk) contributions. gate_seq_idx = torch.arange(seqlen, device=q.device, - dtype=torch.int32).unsqueeze(0).expand( - num_filtered_chunk, seqlen) + dtype=torch.int32).unsqueeze(0).expand(num_filtered_chunk, seqlen) chunk_start = cu_chunk[filtered_chunk_indices] # [num_filtered_chunk] chunk_end = cu_chunk[filtered_chunk_indices + 1] # [num_filtered_chunk] - gate_self_chunk_mask = ( - (gate_seq_idx >= chunk_start.unsqueeze(1)) & - (gate_seq_idx < chunk_end.unsqueeze(1))).unsqueeze(1).expand( - -1, num_head, -1) + gate_self_chunk_mask = ((gate_seq_idx >= chunk_start.unsqueeze(1)) & + (gate_seq_idx < chunk_end.unsqueeze(1))).unsqueeze(1).expand(-1, num_head, -1) amplification_factor = 1e9 # Example factor; adjust as needed. origin_gate = gate.clone() gate = gate.clone() @@ -664,11 +592,7 @@ def moba_attn_varlen( if threshold_type == 'query_head': # === per‐ top-k across chunks (original behavior) === # gate: (C, H, S) - _, gate_topk_idx = torch.topk(gate, - k=moba_topk, - dim=0, - largest=True, - sorted=False) + _, gate_topk_idx = torch.topk(gate, k=moba_topk, dim=0, largest=True, sorted=False) gate_idx_mask = torch.zeros_like(gate, dtype=torch.bool) gate_idx_mask.scatter_(0, gate_topk_idx, True) gate_mask = valid_gate_mask & gate_idx_mask @@ -679,85 +603,52 @@ def moba_attn_varlen( flat_mask = valid_gate_mask.flatten() flat_gate_masked = torch.where(flat_mask, flat_gate, -float("inf")) # pick topk global entries - vals, idx = torch.topk(flat_gate_masked, - k=moba_topk * H * S, - largest=True, - sorted=False) + vals, idx = torch.topk(flat_gate_masked, k=moba_topk * H * S, largest=True, sorted=False) others_mask_flat = torch.zeros_like(flat_mask, dtype=torch.bool) others_mask_flat[idx] = True - gate_mask = (valid_gate_mask.flatten() & others_mask_flat).view( - gate.shape) + gate_mask = (valid_gate_mask.flatten() & others_mask_flat).view(gate.shape) elif threshold_type == 'head_global': # per-head top-k across all chunks and sequence positions C, H, S = gate.shape CS = C * S flat_gate = gate.permute(1, 0, 2).reshape(H, CS) flat_valid = valid_gate_mask.permute(1, 0, 2).reshape(H, CS) - flat_gate_masked = torch.where( - flat_valid, flat_gate, torch.full_like(flat_gate, - -float('inf'))) + flat_gate_masked = torch.where(flat_valid, flat_gate, torch.full_like(flat_gate, -float('inf'))) # pick top-k indices per head - _, topk_idx = torch.topk(flat_gate_masked, - k=moba_topk * S, - dim=1, - largest=True, - sorted=False) + _, topk_idx = torch.topk(flat_gate_masked, k=moba_topk * S, dim=1, largest=True, sorted=False) gate_idx_flat = torch.zeros_like(flat_valid, dtype=torch.bool) gate_idx_flat.scatter_(1, topk_idx, True) gate_mask = gate_idx_flat.reshape(H, C, S).permute(1, 0, 2) else: - raise ValueError( - f"Invalid threshold_type for topk: {threshold_type}. " - "Choose 'query_head', 'block', or 'overall'.") + raise ValueError(f"Invalid threshold_type for topk: {threshold_type}. " + "Choose 'query_head', 'block', or 'overall'.") elif select_mode == 'threshold': # Delegate to the specific thresholding function valid_gate_mask = gate != -float("inf") # (num_chunk, num_head, seqlen) if threshold_type == 'query_head': - gate_mask = _select_threshold_query_head(gate, valid_gate_mask, - gate_self_chunk_mask, - simsum_threshold) + gate_mask = _select_threshold_query_head(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold) elif threshold_type == 'block': - gate_mask = _select_threshold_block(gate, valid_gate_mask, - gate_self_chunk_mask, - simsum_threshold) + gate_mask = _select_threshold_block(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold) elif threshold_type == 'overall': - gate_mask = _select_threshold_overall(gate, valid_gate_mask, - gate_self_chunk_mask, - simsum_threshold) + gate_mask = _select_threshold_overall(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold) elif threshold_type == 'head_global': - gate_mask = _select_threshold_head_global(gate, valid_gate_mask, - gate_self_chunk_mask, - simsum_threshold) + gate_mask = _select_threshold_head_global(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold) else: - raise ValueError( - f"Invalid threshold_type: {threshold_type}. Choose 'query_head', 'block', or 'overall'." - ) + raise ValueError(f"Invalid threshold_type: {threshold_type}. Choose 'query_head', 'block', or 'overall'.") else: - raise ValueError( - f"Invalid select_mode: {select_mode}. Choose 'topk' or 'threshold'." - ) + raise ValueError(f"Invalid select_mode: {select_mode}. Choose 'topk' or 'threshold'.") # eliminate self_chunk in MoBA branch gate_mask = gate_mask & ~gate_self_chunk_mask # if gate_mask is all false, perform flash_attn instead if gate_mask.sum() == 0: - return flash_attn_varlen_func(q, - k, - v, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - causal=False) + return flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, causal=False) # Determine which query positions are selected. # nonzero_indices has shape [N, 3] where each row is [chunk_index, head_index, seq_index]. - moba_q_indices = gate_mask.reshape( - gate_mask.shape[0], -1).nonzero(as_tuple=True)[-1] # [(h s k)] - moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + ( - moba_q_indices // seqlen) - moba_q = rearrange(q, "s h d -> (h s) d").index_select( - 0, moba_q_indices).unsqueeze(1) + moba_q_indices = gate_mask.reshape(gate_mask.shape[0], -1).nonzero(as_tuple=True)[-1] # [(h s k)] + moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + (moba_q_indices // seqlen) + moba_q = rearrange(q, "s h d -> (h s) d").index_select(0, moba_q_indices).unsqueeze(1) # Build cumulative sequence lengths for the selected queries. moba_seqlen_q = gate_mask.sum(dim=-1).flatten() @@ -775,23 +666,19 @@ def moba_attn_varlen( # Rearrange gathered KV for the MOBA branch. experts_tensor = rearrange(gathered, "nc cl two h d -> (nc h) cl two d") - valid_expert_lengths = chunk_lengths.unsqueeze(1).expand( - num_filtered_chunk, num_head).reshape(-1).to(torch.int32) + valid_expert_lengths = chunk_lengths.unsqueeze(1).expand(num_filtered_chunk, num_head).reshape(-1).to(torch.int32) if q_zero_mask.sum() > 0: experts_tensor = experts_tensor[valid_expert_mask] valid_expert_lengths = valid_expert_lengths[valid_expert_mask] - seq_range = torch.arange(experts_tensor.shape[1], - device=experts_tensor.device).unsqueeze(0) + seq_range = torch.arange(experts_tensor.shape[1], device=experts_tensor.device).unsqueeze(0) mask = seq_range < valid_expert_lengths.unsqueeze(1) moba_kv = experts_tensor[mask] # Shape: ((nc h cl_valid) two d) moba_kv = moba_kv.unsqueeze(2) # Shape: ((nc h cl_valid) two 1 d) moba_cu_seqlen_kv = torch.cat( - [ - torch.zeros(1, device=experts_tensor.device, dtype=torch.int32), - valid_expert_lengths.cumsum(dim=0) - ], + [torch.zeros(1, device=experts_tensor.device, dtype=torch.int32), + valid_expert_lengths.cumsum(dim=0)], dim=0, ).to(torch.int32) @@ -831,18 +718,14 @@ def process_moba_input( torch.Tensor: Processed input tensor. """ if isinstance(chunk_size, float) or isinstance(chunk_size, int): - moba_chunk_size = int(chunk_size * patch_resolution[1] * - patch_resolution[2]) + moba_chunk_size = int(chunk_size * patch_resolution[1] * patch_resolution[2]) else: - assert isinstance( - chunk_size, (tuple, list) - ), f"chunk_size should be a tuple, list, or int, now it is: {type(chunk_size)}" + assert isinstance(chunk_size, + (tuple, list)), f"chunk_size should be a tuple, list, or int, now it is: {type(chunk_size)}" if len(chunk_size) == 2: - assert patch_resolution[1] % chunk_size[ - 0] == 0 and patch_resolution[2] % chunk_size[ - 1] == 0, f"spatial patch_resolution {patch_resolution[1:]} should be divisible by 2d chunk_size {chunk_size}" - nch, ncw = patch_resolution[1] // chunk_size[0], patch_resolution[ - 2] // chunk_size[1] + assert patch_resolution[1] % chunk_size[0] == 0 and patch_resolution[2] % chunk_size[ + 1] == 0, f"spatial patch_resolution {patch_resolution[1:]} should be divisible by 2d chunk_size {chunk_size}" + nch, ncw = patch_resolution[1] // chunk_size[0], patch_resolution[2] // chunk_size[1] x = rearrange(x, "b (t nch ch ncw cw) n d -> b (nch ncw t ch cw) n d", t=patch_resolution[0], @@ -852,26 +735,22 @@ def process_moba_input( cw=chunk_size[1]) moba_chunk_size = patch_resolution[0] * chunk_size[0] * chunk_size[1] elif len(chunk_size) == 3: - assert patch_resolution[0] % chunk_size[0] == 0 and patch_resolution[ - 1] % chunk_size[1] == 0 and patch_resolution[2] % chunk_size[ + assert patch_resolution[0] % chunk_size[0] == 0 and patch_resolution[1] % chunk_size[ + 1] == 0 and patch_resolution[2] % chunk_size[ 2] == 0, f"patch_resolution {patch_resolution} should be divisible by 3d chunk_size {chunk_size}" - nct, nch, ncw = patch_resolution[0] // chunk_size[ - 0], patch_resolution[1] // chunk_size[1], patch_resolution[ - 2] // chunk_size[2] - x = rearrange( - x, - "b (nct ct nch ch ncw cw) n d -> b (nct nch ncw ct ch cw) n d", - nct=nct, - nch=nch, - ncw=ncw, - ct=chunk_size[0], - ch=chunk_size[1], - cw=chunk_size[2]) + nct, nch, ncw = patch_resolution[0] // chunk_size[0], patch_resolution[1] // chunk_size[ + 1], patch_resolution[2] // chunk_size[2] + x = rearrange(x, + "b (nct ct nch ch ncw cw) n d -> b (nct nch ncw ct ch cw) n d", + nct=nct, + nch=nch, + ncw=ncw, + ct=chunk_size[0], + ch=chunk_size[1], + cw=chunk_size[2]) moba_chunk_size = chunk_size[0] * chunk_size[1] * chunk_size[2] else: - raise ValueError( - f"chunk_size should be a int, or a tuple of length 2 or 3, now it is: {len(chunk_size)}" - ) + raise ValueError(f"chunk_size should be a int, or a tuple of length 2 or 3, now it is: {len(chunk_size)}") return x, moba_chunk_size @@ -892,15 +771,14 @@ def process_moba_output( ch=chunk_size[0], cw=chunk_size[1]) elif len(chunk_size) == 3: - x = rearrange( - x, - "b (nct nch ncw ct ch cw) n d -> b (nct ct nch ch ncw cw) n d", - nct=patch_resolution[0] // chunk_size[0], - nch=patch_resolution[1] // chunk_size[1], - ncw=patch_resolution[2] // chunk_size[2], - ct=chunk_size[0], - ch=chunk_size[1], - cw=chunk_size[2]) + x = rearrange(x, + "b (nct nch ncw ct ch cw) n d -> b (nct ct nch ch ncw cw) n d", + nct=patch_resolution[0] // chunk_size[0], + nch=patch_resolution[1] // chunk_size[1], + ncw=patch_resolution[2] // chunk_size[2], + ct=chunk_size[0], + ch=chunk_size[1], + cw=chunk_size[2]) return x @@ -912,18 +790,11 @@ def generate_data(batch_size, seqlen, num_head, head_dim, dtype): torch.cuda.manual_seed(0) device = torch.cuda.current_device() - q = torch.randn((batch_size, seqlen, num_head, head_dim), - requires_grad=True).to(dtype=dtype, device='cuda') - k = torch.randn((batch_size, seqlen, num_head, head_dim), - requires_grad=True).to(dtype=dtype, device='cuda') - v = torch.randn((batch_size, seqlen, num_head, head_dim), - requires_grad=True).to(dtype=dtype, device='cuda') + q = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda') + k = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda') + v = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda') print(f"q.shape: {q.shape}, k.shape: {k.shape}, v.shape: {v.shape}") - cu_seqlens = torch.arange(0, - q.shape[0] * q.shape[1] + 1, - q.shape[1], - dtype=torch.int32, - device='cuda') + cu_seqlens = torch.arange(0, q.shape[0] * q.shape[1] + 1, q.shape[1], dtype=torch.int32, device='cuda') max_seqlen = q.shape[1] q = rearrange(q, "b s ... -> (b s) ...") k = rearrange(k, "b s ... -> (b s) ...") @@ -944,8 +815,7 @@ def test_attn_varlen_moba_speed(batch, threshold_type='query_head'): """Speed test comparing flash_attn vs moba_attention""" # Get data - q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head, - head_dim, dtype) + q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head, head_dim, dtype) print( f"batch:{batch} head:{head} seqlen:{seqlen} chunk:{moba_chunk_size} topk:{moba_topk} select_mode: {select_mode} simsum_threshold:{simsum_threshold}" ) @@ -957,27 +827,13 @@ def test_attn_varlen_moba_speed(batch, # Warmup for _ in range(warmup_iters): - o = flash_attn_varlen_func(q, - k, - v, - cu_seqlen, - cu_seqlen, - max_seqlen, - max_seqlen, - causal=False) + o = flash_attn_varlen_func(q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False) torch.autograd.backward(o, vo_grad) torch.cuda.synchronize() start_flash = time.perf_counter() for _ in range(perf_test_iters): - o = flash_attn_varlen_func(q, - k, - v, - cu_seqlen, - cu_seqlen, - max_seqlen, - max_seqlen, - causal=False) + o = flash_attn_varlen_func(q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False) torch.autograd.backward(o, vo_grad) torch.cuda.synchronize() diff --git a/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py b/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py index d0dc1b840..d0d647cfe 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py +++ b/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py @@ -78,7 +78,7 @@ def construct_variable_block_sizes( n_t, n_h, n_w = num_tiles def _sizes(dim_len: int, tile: int, n: int) -> torch.LongTensor: - sizes = torch.full((n,), tile, dtype=torch.int, device=device) + sizes = torch.full((n, ), tile, dtype=torch.int, device=device) remainder = dim_len - (n - 1) * tile sizes[-1] = remainder if remainder > 0 else tile return sizes @@ -87,11 +87,7 @@ def construct_variable_block_sizes( h_sizes = _sizes(h, ts_h, n_h) w_sizes = _sizes(w, ts_w, n_w) - return ( - t_sizes[:, None, None] - * h_sizes[None, :, None] - * w_sizes[None, None, :] - ).reshape(-1) + return (t_sizes[:, None, None] * h_sizes[None, :, None] * w_sizes[None, None, :]).reshape(-1) def get_non_pad_index( @@ -134,10 +130,8 @@ def build_vsa_metadata( ts_t, ts_h, ts_w = tile_size max_block_size = math.prod(tile_size) if max_block_size not in _SUPPORTED_VSA_BLOCK_VOLUMES: - raise ValueError( - f"Unsupported VSA tile volume {max_block_size} for tile_size={tile_size}; " - f"supported volumes are {_SUPPORTED_VSA_BLOCK_VOLUMES}." - ) + raise ValueError(f"Unsupported VSA tile volume {max_block_size} for tile_size={tile_size}; " + f"supported volumes are {_SUPPORTED_VSA_BLOCK_VOLUMES}.") num_tiles = ( math.ceil(T / ts_t), diff --git a/fastvideo-kernel/tests/test_attn_qat_infer.py b/fastvideo-kernel/tests/test_attn_qat_infer.py index 270b4d201..b231cbc1c 100644 --- a/fastvideo-kernel/tests/test_attn_qat_infer.py +++ b/fastvideo-kernel/tests/test_attn_qat_infer.py @@ -16,6 +16,7 @@ or via pytest: import os import sys + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) import pytest @@ -42,13 +43,10 @@ def cosine_similarity(a: torch.Tensor, b: torch.Tensor) -> float: return (torch.dot(a_flat, b_flat) / (norm_a * norm_b)).item() -def reference_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - is_causal: bool) -> torch.Tensor: +def reference_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, is_causal: bool) -> torch.Tensor: """PyTorch math-backend SDPA reference in float32.""" with sdpa_kernel([SDPBackend.MATH]): - out = F.scaled_dot_product_attention( - q.float(), k.float(), v.float(), is_causal=is_causal - ) + out = F.scaled_dot_product_attention(q.float(), k.float(), v.float(), is_causal=is_causal) return out.to(q.dtype) @@ -56,6 +54,7 @@ def reference_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, # Benchmark # --------------------------------------------------------------------------- + def bench(B, H, L, D, is_causal, warmup=20, iters=100): q = torch.randn(B, H, L, D, device="cuda", dtype=torch.bfloat16) k = torch.randn(B, H, L, D, device="cuda", dtype=torch.bfloat16) @@ -87,10 +86,8 @@ def run_benchmark(): for L in [512, 1024, 2048, 4096]: for D in [64, 128]: ms, tflops = bench(B, H, L, D, is_causal) - print( - f"{B:>2} {H:>2} {L:>5} {D:>4} {str(is_causal):>6}" - f" {ms:>8.3f} {tflops:>7.2f}" - ) + print(f"{B:>2} {H:>2} {L:>5} {D:>4} {str(is_causal):>6}" + f" {ms:>8.3f} {tflops:>7.2f}") # --------------------------------------------------------------------------- @@ -101,15 +98,21 @@ import pytest @pytest.mark.parametrize("B,H,L,D", [ - (2, 4, 128, 64), (2, 4, 128, 128), - (2, 4, 256, 64), (2, 4, 256, 128), - (2, 4, 512, 64), (2, 4, 512, 128), - (2, 4, 1024, 64), (2, 4, 1024, 128), - (2, 4, 2048, 64), (2, 4, 2048, 128), - (2, 4, 4096, 64), (2, 4, 4096, 128), - (1, 4, 128, 128), - (2, 8, 256, 128), - (1, 4, 384, 128), + (2, 4, 128, 64), + (2, 4, 128, 128), + (2, 4, 256, 64), + (2, 4, 256, 128), + (2, 4, 512, 64), + (2, 4, 512, 128), + (2, 4, 1024, 64), + (2, 4, 1024, 128), + (2, 4, 2048, 64), + (2, 4, 2048, 128), + (2, 4, 4096, 64), + (2, 4, 4096, 128), + (1, 4, 128, 128), + (2, 8, 256, 128), + (1, 4, 384, 128), ]) @pytest.mark.parametrize("causal", [False, True], ids=["non_causal", "causal"]) def test_accuracy_sdpa(causal: bool, B: int, H: int, L: int, D: int): @@ -135,15 +138,21 @@ if __name__ == "__main__": print("Accuracy sweep — SDPA reference:") for causal in [False, True]: for B, H, L, D in [ - (2, 4, 128, 64), (2, 4, 128, 128), - (2, 4, 256, 64), (2, 4, 256, 128), - (2, 4, 512, 64), (2, 4, 512, 128), - (2, 4, 1024, 64), (2, 4, 1024, 128), - (2, 4, 2048, 64), (2, 4, 2048, 128), - (2, 4, 4096, 64), (2, 4, 4096, 128), - (1, 4, 128, 128), - (2, 8, 256, 128), - (1, 4, 384, 128), + (2, 4, 128, 64), + (2, 4, 128, 128), + (2, 4, 256, 64), + (2, 4, 256, 128), + (2, 4, 512, 64), + (2, 4, 512, 128), + (2, 4, 1024, 64), + (2, 4, 1024, 128), + (2, 4, 2048, 64), + (2, 4, 2048, 128), + (2, 4, 4096, 64), + (2, 4, 4096, 128), + (1, 4, 128, 128), + (2, 8, 256, 128), + (1, 4, 384, 128), ]: test_accuracy_sdpa(causal, B, H, L, D) print() diff --git a/fastvideo-kernel/tests/test_attn_qat_train.py b/fastvideo-kernel/tests/test_attn_qat_train.py index 20783f306..22755e832 100644 --- a/fastvideo-kernel/tests/test_attn_qat_train.py +++ b/fastvideo-kernel/tests/test_attn_qat_train.py @@ -30,15 +30,15 @@ def attn_qat_train_wrapper(q_BLHD, k_BLHD, v_BLHD, is_causal=False, sm_scale=Non """ if sm_scale is None: sm_scale = 1.0 / sqrt(q_BLHD.shape[-1]) - + # Convert from BLHD to BHLD format q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() - + # Call attention with BHLD format o_BHLD = attention(q_BHLD, k_BHLD, v_BHLD, is_causal, sm_scale) - + # Convert back from BHLD to BLHD format return o_BHLD.permute(0, 2, 1, 3).contiguous() @@ -61,16 +61,16 @@ def fused_attn_wrapper(q_BLHD, k_BLHD, v_BLHD, is_causal=False, sm_scale=None, w """ if sm_scale is None: sm_scale = 1.0 / sqrt(q_BLHD.shape[-1]) - + # Convert from BLHD to BHLD format q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() - + # Call fused attention with BHLD format # Note: fused_attention expects inputs in BHLD format and only supports same sequence lengths o_BHLD = fused_attention(q_BHLD, k_BHLD, v_BHLD, is_causal, sm_scale, warp_specialize) - + # Convert back from BHLD to BLHD format return o_BHLD.permute(0, 2, 1, 3).contiguous() @@ -89,16 +89,16 @@ def cosine_similarity(tensor1, tensor2): # Flatten tensors for computation t1_flat = tensor1.flatten().float() t2_flat = tensor2.flatten().float() - + # Compute cosine similarity: (A · B) / (||A|| * ||B||) dot_product = torch.dot(t1_flat, t2_flat) norm1 = torch.norm(t1_flat) norm2 = torch.norm(t2_flat) - + # Avoid division by zero if norm1 == 0 or norm2 == 0: return 0.0 - + cos_sim = dot_product / (norm1 * norm2) return cos_sim.item() @@ -119,7 +119,7 @@ def naive_attention(q, k, v, causal, sm_scale): """ # Compute attention scores: QK^T scores = torch.matmul(q, k.transpose(-2, -1)) * sm_scale - + # Apply causal mask if needed # Note: Causal masking is only meaningful when Q and KV have the same sequence length if causal: @@ -133,98 +133,101 @@ def naive_attention(q, k, v, causal, sm_scale): # For different sequence lengths, causal masking doesn't apply in the same way # In cross-attention, typically no causal masking is applied pass - + # Apply softmax attn_weights = torch.softmax(scores, dim=-1) - + # Apply attention weights to values out = torch.matmul(attn_weights, v) - + return out + def test_qat_attention_forward(): """Test forward pass with QAT enabled vs disabled.""" torch.manual_seed(42) - + # Test parameters Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 - dtype = torch.bfloat16 + dtype = torch.bfloat16 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = False - + # Create input tensors in BLHD format (B, L, H, D) = (Z, N_CTX, H, HEAD_DIM) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) - + # Test with QAT (using wrapper that does permute/contiguous) out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Convert to BHLD for naive attention comparison q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + # Check that outputs have correct shape assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) - + # Check that outputs are finite assert torch.isfinite(out_qat_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() - + # Compare QAT and naive outputs max_diff = (out_qat_BLHD - out_naive_BLHD).abs().max() mean_diff = (out_qat_BLHD - out_naive_BLHD).abs().mean() cos_sim = cosine_similarity(out_qat_BLHD, out_naive_BLHD) - print(f" QAT vs Naive Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM} - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" QAT vs Naive Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM} - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # QAT output should be reasonably close to naive (within tolerance due to quantization) # Using a reasonable tolerance for float16 and quantization effects # assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close, got max_diff={max_diff.item():.6f}" - + print("✓ Forward pass test passed.") def test_qat_attention_backward(): """Test backward pass with QAT enabled.""" torch.manual_seed(42) - + # Test parameters Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 dtype = torch.bfloat16 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - + # Create input tensors in BLHD format for QAT q_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) k_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) v_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) - + # Create input tensors for naive (same values, convert to BHLD) q_naive_BHLD = q_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) k_naive_BHLD = k_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) v_naive_BHLD = v_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) - + # Forward pass with QAT (using wrapper) out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) - + # Forward pass with naive out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + # Create dummy gradient (same for both, in BLHD format) # Ensure gradient is contiguous as required by the backward function dout = torch.randn_like(out_qat_BLHD).contiguous() - + # Backward pass for QAT out_qat_BLHD.backward(dout) - + # Backward pass for naive out_naive_BLHD.backward(dout) - + # Check that gradients are computed assert q_qat_BLHD.grad is not None assert k_qat_BLHD.grad is not None @@ -232,319 +235,339 @@ def test_qat_attention_backward(): assert q_naive_BHLD.grad is not None assert k_naive_BHLD.grad is not None assert v_naive_BHLD.grad is not None - + # Check that gradients have correct shape assert q_qat_BLHD.grad.shape == q_qat_BLHD.shape assert k_qat_BLHD.grad.shape == k_qat_BLHD.shape assert v_qat_BLHD.grad.shape == v_qat_BLHD.shape - + # Check that gradients are finite assert torch.isfinite(q_qat_BLHD.grad).all() assert torch.isfinite(k_qat_BLHD.grad).all() assert torch.isfinite(v_qat_BLHD.grad).all() - + # Convert naive gradients to BLHD format for comparison q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() - + # Compare gradients with naive dq_diff = (q_qat_BLHD.grad - q_naive_grad_BLHD).abs() dk_diff = (k_qat_BLHD.grad - k_naive_grad_BLHD).abs() dv_diff = (v_qat_BLHD.grad - v_naive_grad_BLHD).abs() - + dq_cos_sim = cosine_similarity(q_qat_BLHD.grad, q_naive_grad_BLHD) dk_cos_sim = cosine_similarity(k_qat_BLHD.grad, k_naive_grad_BLHD) dv_cos_sim = cosine_similarity(v_qat_BLHD.grad, v_naive_grad_BLHD) - - print(f" dQ - Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}") - print(f" dK - Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}") - print(f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}") - + + print( + f" dQ - Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}" + ) + print( + f" dK - Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}" + ) + print( + f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}" + ) + # Gradients should be reasonably close (within tolerance due to quantization) # assert dq_diff.max() < 2.0, f"dQ gradients should be reasonably close, got max_diff={dq_diff.max().item():.6f}" # assert dk_diff.max() < 2.0, f"dK gradients should be reasonably close, got max_diff={dk_diff.max().item():.6f}" # assert dv_diff.max() < 2.0, f"dV gradients should be reasonably close, got max_diff={dv_diff.max().item():.6f}" - + print("✓ Backward pass test passed.") def test_qat_attention_different_shapes(): """Test QAT attention with different input shapes.""" torch.manual_seed(42) - + test_configs = [ (2, 4, 128, 64), # Medium - (1, 8, 256, 128), # Large head dim + (1, 8, 256, 128), # Large head dim ] - + dtype = torch.bfloat16 causal = True - + for Z, H, N_CTX, HEAD_DIM in test_configs: # Create input tensors in BLHD format q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) - + sm_scale = 1.0 / sqrt(HEAD_DIM) out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert torch.isfinite(out_qat_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() - + # Compare outputs max_diff = (out_qat_BLHD - out_naive_BLHD).abs().max() mean_diff = (out_qat_BLHD - out_naive_BLHD).abs().mean() cos_sim = cosine_similarity(out_qat_BLHD, out_naive_BLHD) - print(f" (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # Outputs should be reasonably close # assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close for shape (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}), got max_diff={max_diff.item():.6f}" - + print(f"✓ Shape test passed for (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM})") def test_qat_attention_non_causal(): """Test QAT attention with non-causal masking.""" torch.manual_seed(42) - + Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 dtype = torch.bfloat16 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = False - + # Create input tensors in BLHD format q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) - + out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert torch.isfinite(out_qat_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() - + # Compare outputs max_diff = (out_qat_BLHD - out_naive_BLHD).abs().max() mean_diff = (out_qat_BLHD - out_naive_BLHD).abs().mean() cos_sim = cosine_similarity(out_qat_BLHD, out_naive_BLHD) - print(f" QAT vs Naive (non-causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" QAT vs Naive (non-causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # Outputs should be reasonably close # assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close for non-causal, got max_diff={max_diff.item():.6f}" - + print("✓ Non-causal test passed.") def test_qat_attention_causal(): """Test QAT attention with causal masking.""" torch.manual_seed(42) - + Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 dtype = torch.bfloat16 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - + # Create input tensors in BLHD format q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) - + out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert torch.isfinite(out_qat_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() - + # Compare outputs max_diff = (out_qat_BLHD - out_naive_BLHD).abs().max() mean_diff = (out_qat_BLHD - out_naive_BLHD).abs().mean() cos_sim = cosine_similarity(out_qat_BLHD, out_naive_BLHD) - print(f" QAT vs Naive (causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" QAT vs Naive (causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # Outputs should be reasonably close # assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close for causal, got max_diff={max_diff.item():.6f}" - + print("✓ Causal test passed.") def test_qat_attention_causal_backward(): """Test backward pass of QAT attention with causal masking vs naive.""" torch.manual_seed(42) - + Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 dtype = torch.bfloat16 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - + # Create input tensors in BLHD format for QAT q_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) k_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) v_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) - + # Create input tensors for naive (same values, convert to BHLD) q_naive_BHLD = q_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) k_naive_BHLD = k_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) v_naive_BHLD = v_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) - + # Forward pass with QAT out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) - + # Forward pass with naive out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale) out_naive_BHLD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + # Create dummy gradient dout = torch.randn_like(out_qat_BLHD).contiguous() - + # Backward pass for QAT out_qat_BLHD.backward(dout) - + # Backward pass for naive out_naive_BHLD.backward(dout) - + # Check that gradients are computed assert q_qat_BLHD.grad is not None assert k_qat_BLHD.grad is not None assert v_qat_BLHD.grad is not None - + # Convert naive gradients to BLHD format for comparison q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() - + # Compare gradients dq_diff = (q_qat_BLHD.grad - q_naive_grad_BLHD).abs() dk_diff = (k_qat_BLHD.grad - k_naive_grad_BLHD).abs() dv_diff = (v_qat_BLHD.grad - v_naive_grad_BLHD).abs() - + dq_cos_sim = cosine_similarity(q_qat_BLHD.grad, q_naive_grad_BLHD) dk_cos_sim = cosine_similarity(k_qat_BLHD.grad, k_naive_grad_BLHD) dv_cos_sim = cosine_similarity(v_qat_BLHD.grad, v_naive_grad_BLHD) - - print(f" dQ - Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}") - print(f" dK - Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}") - print(f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}") - + + print( + f" dQ - Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}" + ) + print( + f" dK - Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}" + ) + print( + f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}" + ) + print("✓ Causal backward test passed.") def test_qat_attention_different_seq_lengths(): """Test QAT attention with different sequence lengths for Q and KV (cross-attention).""" torch.manual_seed(42) - + test_configs = [ - (2, 4, 64, 128, 64, False), # Q shorter than KV, non-causal - (2, 4, 128, 64, 64, False), # Q longer than KV, non-causal - (1, 8, 256, 128, 64, False), # Q longer than KV, larger dimensions + (2, 4, 64, 128, 64, False), # Q shorter than KV, non-causal + (2, 4, 128, 64, 64, False), # Q longer than KV, non-causal + (1, 8, 256, 128, 64, False), # Q longer than KV, larger dimensions ] - + dtype = torch.bfloat16 - + for Z, H, N_CTX_Q, N_CTX_KV, HEAD_DIM, causal in test_configs: sm_scale = 1.0 / sqrt(HEAD_DIM) - + # Create input tensors in BLHD format with different sequence lengths q_BLHD = torch.randn((Z, N_CTX_Q, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE) - + # Test with QAT (using wrapper that does permute/contiguous) out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Convert to BHLD for naive attention comparison q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + # Check that outputs have correct shape (should match Q sequence length) assert out_qat_BLHD.shape == (Z, N_CTX_Q, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX_Q, H, HEAD_DIM) - + # Check that outputs are finite assert torch.isfinite(out_qat_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() - + # Compare QAT and naive outputs max_diff = (out_qat_BLHD - out_naive_BLHD).abs().max() mean_diff = (out_qat_BLHD - out_naive_BLHD).abs().mean() cos_sim = cosine_similarity(out_qat_BLHD, out_naive_BLHD) - print(f" (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}, causal={causal}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}, causal={causal}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # Outputs should be reasonably close # assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close for (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}), got max_diff={max_diff.item():.6f}" - + print(f"✓ Different sequence lengths test passed for (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV})") def test_qat_attention_different_seq_lengths_backward(): """Test backward pass of QAT attention with different sequence lengths for Q and KV (cross-attention).""" torch.manual_seed(42) - + test_configs = [ - (2, 4, 64, 128, 64, False), # Q shorter than KV, non-causal - (2, 4, 128, 64, 64, False), # Q longer than KV, non-causal + (2, 4, 64, 128, 64, False), # Q shorter than KV, non-causal + (2, 4, 128, 64, 64, False), # Q longer than KV, non-causal ] - + dtype = torch.bfloat16 - + for Z, H, N_CTX_Q, N_CTX_KV, HEAD_DIM, causal in test_configs: sm_scale = 1.0 / sqrt(HEAD_DIM) - + # Create input tensors in BLHD format with different sequence lengths q_qat_BLHD = torch.randn((Z, N_CTX_Q, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) k_qat_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) v_qat_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) - + # Create input tensors for naive (same values, convert to BHLD) q_naive_BHLD = q_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) k_naive_BHLD = k_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) v_naive_BHLD = v_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) - + # Forward pass with QAT out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) - + # Forward pass with naive out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale) out_naive_BHLD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + # Create dummy gradient dout = torch.randn_like(out_qat_BLHD).contiguous() - + # Backward pass for QAT out_qat_BLHD.backward(dout) - + # Backward pass for naive out_naive_BHLD.backward(dout) - + # Check that gradients are computed assert q_qat_BLHD.grad is not None assert k_qat_BLHD.grad is not None @@ -552,124 +575,132 @@ def test_qat_attention_different_seq_lengths_backward(): assert q_naive_BHLD.grad is not None assert k_naive_BHLD.grad is not None assert v_naive_BHLD.grad is not None - + # Check that gradients have correct shape assert q_qat_BLHD.grad.shape == q_qat_BLHD.shape assert k_qat_BLHD.grad.shape == k_qat_BLHD.shape assert v_qat_BLHD.grad.shape == v_qat_BLHD.shape - + # Check that gradients are finite assert torch.isfinite(q_qat_BLHD.grad).all() assert torch.isfinite(k_qat_BLHD.grad).all() assert torch.isfinite(v_qat_BLHD.grad).all() - + # Convert naive gradients to BLHD format for comparison q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() - + # Compare gradients dq_diff = (q_qat_BLHD.grad - q_naive_grad_BLHD).abs() dk_diff = (k_qat_BLHD.grad - k_naive_grad_BLHD).abs() dv_diff = (v_qat_BLHD.grad - v_naive_grad_BLHD).abs() - + dq_cos_sim = cosine_similarity(q_qat_BLHD.grad, q_naive_grad_BLHD) dk_cos_sim = cosine_similarity(k_qat_BLHD.grad, k_naive_grad_BLHD) dv_cos_sim = cosine_similarity(v_qat_BLHD.grad, v_naive_grad_BLHD) - - print(f" (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}) - dQ: Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}") - print(f" (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}) - dK: Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}") - print(f" (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}) - dV: Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}") - + + print( + f" (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}) - dQ: Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}" + ) + print( + f" (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}) - dK: Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}" + ) + print( + f" (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}) - dV: Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}" + ) + # Gradients should be reasonably close (within tolerance due to quantization) # assert dq_diff.max() < 2.0, f"dQ gradients should be reasonably close, got max_diff={dq_diff.max().item():.6f}" # assert dk_diff.max() < 2.0, f"dK gradients should be reasonably close, got max_diff={dk_diff.max().item():.6f}" # assert dv_diff.max() < 2.0, f"dV gradients should be reasonably close, got max_diff={dv_diff.max().item():.6f}" - + print(f"✓ Cross-attention backward test passed for (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV})") def test_fused_attention_forward(): """Test forward pass of fused attention vs naive PyTorch implementation.""" torch.manual_seed(42) - + # Test parameters Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 dtype = torch.bfloat16 # fused_attention uses float16 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = False - + # Create input tensors in BLHD format (B, L, H, D) = (Z, N_CTX, H, HEAD_DIM) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) - + # Test with fused attention (using wrapper that does permute/contiguous) out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Convert to BHLD for naive attention comparison q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + # Check that outputs have correct shape assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) - + # Check that outputs are finite assert torch.isfinite(out_fused_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() - + # Compare fused and naive outputs max_diff = (out_fused_BLHD - out_naive_BLHD).abs().max() mean_diff = (out_fused_BLHD - out_naive_BLHD).abs().mean() cos_sim = cosine_similarity(out_fused_BLHD, out_naive_BLHD) - print(f" Fused vs Naive - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" Fused vs Naive - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # Fused attention should be reasonably close to naive (within tolerance for float16) # assert max_diff < 1e-1, f"Fused and naive outputs should be reasonably close, got max_diff={max_diff.item():.6f}" - + print("✓ Fused attention forward pass test passed.") def test_fused_attention_backward(): """Test backward pass of fused attention vs naive PyTorch implementation.""" torch.manual_seed(42) - + # Test parameters Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 dtype = torch.bfloat16 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - + # Create input tensors in BLHD format for fused attention q_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) k_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) v_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) - + # Create input tensors for naive (same values, convert to BHLD) q_naive_BHLD = q_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) k_naive_BHLD = k_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) v_naive_BHLD = v_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) - + # Forward pass with fused attention (using wrapper) out_fused_BLHD = fused_attn_wrapper(q_fused_BLHD, k_fused_BLHD, v_fused_BLHD, causal, sm_scale) - + # Forward pass with naive out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + # Create dummy gradient (same for both, in BLHD format) dout = torch.randn_like(out_fused_BLHD).contiguous() - + # Backward pass for fused attention out_fused_BLHD.backward(dout) - + # Backward pass for naive out_naive_BLHD.backward(dout) - + # Check that gradients are computed assert q_fused_BLHD.grad is not None assert k_fused_BLHD.grad is not None @@ -677,312 +708,332 @@ def test_fused_attention_backward(): assert q_naive_BHLD.grad is not None assert k_naive_BHLD.grad is not None assert v_naive_BHLD.grad is not None - + # Check that gradients have correct shape assert q_fused_BLHD.grad.shape == q_fused_BLHD.shape assert k_fused_BLHD.grad.shape == k_fused_BLHD.shape assert v_fused_BLHD.grad.shape == v_fused_BLHD.shape - + # Check that gradients are finite assert torch.isfinite(q_fused_BLHD.grad).all() assert torch.isfinite(k_fused_BLHD.grad).all() assert torch.isfinite(v_fused_BLHD.grad).all() - + # Convert naive gradients to BLHD format for comparison q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() - + # Compare gradients with naive dq_diff = (q_fused_BLHD.grad - q_naive_grad_BLHD).abs() dk_diff = (k_fused_BLHD.grad - k_naive_grad_BLHD).abs() dv_diff = (v_fused_BLHD.grad - v_naive_grad_BLHD).abs() - + dq_cos_sim = cosine_similarity(q_fused_BLHD.grad, q_naive_grad_BLHD) dk_cos_sim = cosine_similarity(k_fused_BLHD.grad, k_naive_grad_BLHD) dv_cos_sim = cosine_similarity(v_fused_BLHD.grad, v_naive_grad_BLHD) - - print(f" dQ - Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}") - print(f" dK - Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}") - print(f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}") - + + print( + f" dQ - Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}" + ) + print( + f" dK - Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}" + ) + print( + f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}" + ) + # Gradients should be reasonably close (within tolerance for float16) # assert dq_diff.max() < 1e-1, f"dQ gradients should be reasonably close, got max_diff={dq_diff.max().item():.6f}" # assert dk_diff.max() < 1e-1, f"dK gradients should be reasonably close, got max_diff={dk_diff.max().item():.6f}" # assert dv_diff.max() < 1e-1, f"dV gradients should be reasonably close, got max_diff={dv_diff.max().item():.6f}" - + print("✓ Fused attention backward pass test passed.") def test_fused_attention_different_shapes(): """Test fused attention with different input shapes.""" torch.manual_seed(42) - + test_configs = [ - (1, 2, 64, 32), # Small + (1, 2, 64, 32), # Small (2, 4, 128, 64), # Medium - (1, 8, 256, 128), # Large head dim + (1, 8, 256, 128), # Large head dim ] - + dtype = torch.bfloat16 causal = True - + for Z, H, N_CTX, HEAD_DIM in test_configs: # Create input tensors in BLHD format q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) - + sm_scale = 1.0 / sqrt(HEAD_DIM) out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert torch.isfinite(out_fused_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() - + # Compare outputs max_diff = (out_fused_BLHD - out_naive_BLHD).abs().max() mean_diff = (out_fused_BLHD - out_naive_BLHD).abs().mean() cos_sim = cosine_similarity(out_fused_BLHD, out_naive_BLHD) - print(f" (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # Outputs should be reasonably close # assert max_diff < 1e-1, f"Fused and naive outputs should be reasonably close for shape (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}), got max_diff={max_diff.item():.6f}" - + print(f"✓ Fused attention shape test passed for (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM})") def test_fused_attention_non_causal(): """Test fused attention with non-causal masking.""" torch.manual_seed(42) - + Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 dtype = torch.bfloat16 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = False - + # Create input tensors in BLHD format q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) - + out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert torch.isfinite(out_fused_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() - + # Compare outputs max_diff = (out_fused_BLHD - out_naive_BLHD).abs().max() mean_diff = (out_fused_BLHD - out_naive_BLHD).abs().mean() cos_sim = cosine_similarity(out_fused_BLHD, out_naive_BLHD) - print(f" Fused vs Naive (non-causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" Fused vs Naive (non-causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # Outputs should be reasonably close # assert max_diff < 1e-1, f"Fused and naive outputs should be reasonably close for non-causal, got max_diff={max_diff.item():.6f}" - + print("✓ Fused attention non-causal test passed.") def test_fused_attention_causal(): """Test fused attention with causal masking.""" torch.manual_seed(42) - + Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 dtype = torch.bfloat16 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - + # Create input tensors in BLHD format q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) - + out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert torch.isfinite(out_fused_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() - + # Compare outputs max_diff = (out_fused_BLHD - out_naive_BLHD).abs().max() mean_diff = (out_fused_BLHD - out_naive_BLHD).abs().mean() cos_sim = cosine_similarity(out_fused_BLHD, out_naive_BLHD) - print(f" Fused vs Naive (causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" Fused vs Naive (causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # Outputs should be reasonably close # assert max_diff < 1e-1, f"Fused and naive outputs should be reasonably close for causal, got max_diff={max_diff.item():.6f}" - + print("✓ Fused attention causal test passed.") def test_fused_attention_causal_backward(): """Test backward pass of fused attention with causal masking vs naive.""" torch.manual_seed(42) - + Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 dtype = torch.bfloat16 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - + # Create input tensors in BLHD format for fused attention q_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) k_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) v_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) - + # Create input tensors for naive (same values, convert to BHLD) q_naive_BHLD = q_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) k_naive_BHLD = k_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) v_naive_BHLD = v_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) - + # Forward pass with fused attention out_fused_BLHD = fused_attn_wrapper(q_fused_BLHD, k_fused_BLHD, v_fused_BLHD, causal, sm_scale) - + # Forward pass with naive out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale) out_naive_BHLD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - + # Create dummy gradient dout = torch.randn_like(out_fused_BLHD).contiguous() - + # Backward pass for fused attention out_fused_BLHD.backward(dout) - + # Backward pass for naive out_naive_BHLD.backward(dout) - + # Check that gradients are computed assert q_fused_BLHD.grad is not None assert k_fused_BLHD.grad is not None assert v_fused_BLHD.grad is not None - + # Convert naive gradients to BLHD format for comparison q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() - + # Compare gradients dq_diff = (q_fused_BLHD.grad - q_naive_grad_BLHD).abs() dk_diff = (k_fused_BLHD.grad - k_naive_grad_BLHD).abs() dv_diff = (v_fused_BLHD.grad - v_naive_grad_BLHD).abs() - + dq_cos_sim = cosine_similarity(q_fused_BLHD.grad, q_naive_grad_BLHD) dk_cos_sim = cosine_similarity(k_fused_BLHD.grad, k_naive_grad_BLHD) dv_cos_sim = cosine_similarity(v_fused_BLHD.grad, v_naive_grad_BLHD) - - print(f" dQ - Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}") - print(f" dK - Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}") - print(f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}") - + + print( + f" dQ - Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}" + ) + print( + f" dK - Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}" + ) + print( + f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}" + ) + print("✓ Fused attention causal backward test passed.") def test_fused_vs_qat_attention_forward(): """Test forward pass comparing fused attention vs QAT attention.""" torch.manual_seed(42) - + # Test parameters Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = False - + # Create input tensors in BLHD format # Use float16 for both (QAT can handle float16, though it typically uses float16) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) - + # Test with fused attention out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Test with QAT attention (convert to float16 for QAT) q_qat_BLHD = q_BLHD.clone().to(torch.bfloat16) k_qat_BLHD = k_BLHD.clone().to(torch.bfloat16) v_qat_BLHD = v_BLHD.clone().to(torch.bfloat16) out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) - + # Convert QAT output back to float16 for comparison out_qat_BLHD = out_qat_BLHD.to(torch.bfloat16) - + # Check that outputs have correct shape assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) - + # Check that outputs are finite assert torch.isfinite(out_fused_BLHD).all() assert torch.isfinite(out_qat_BLHD).all() - + # Compare fused and QAT outputs max_diff = (out_fused_BLHD - out_qat_BLHD).abs().max() mean_diff = (out_fused_BLHD - out_qat_BLHD).abs().mean() cos_sim = cosine_similarity(out_fused_BLHD, out_qat_BLHD) - print(f" Fused vs QAT - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" Fused vs QAT - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # They may differ due to quantization in QAT and different implementations # assert max_diff < 1.0, f"Fused and QAT outputs should be reasonably close, got max_diff={max_diff.item():.6f}" - + print("✓ Fused vs QAT forward pass test passed.") def test_fused_vs_qat_attention_backward(): """Test backward pass comparing fused attention vs QAT attention.""" torch.manual_seed(42) - + # Test parameters Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - + # Create input tensors in BLHD format for fused attention (float16) q_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True) k_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True) v_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True) - + # Create input tensors for QAT (float16, same values) q_qat_BLHD = q_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True) k_qat_BLHD = k_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True) v_qat_BLHD = v_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True) - + # Forward pass with fused attention out_fused_BLHD = fused_attn_wrapper(q_fused_BLHD, k_fused_BLHD, v_fused_BLHD, causal, sm_scale) - + # Forward pass with QAT out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) - + # Create dummy gradient (same for both, in BLHD format) dout = torch.randn_like(out_fused_BLHD).contiguous() dout_qat = dout.to(torch.bfloat16) - + # Backward pass for fused attention out_fused_BLHD.backward(dout) - + # Backward pass for QAT out_qat_BLHD.backward(dout_qat) - + # Check that gradients are computed assert q_fused_BLHD.grad is not None assert k_fused_BLHD.grad is not None @@ -990,203 +1041,215 @@ def test_fused_vs_qat_attention_backward(): assert q_qat_BLHD.grad is not None assert k_qat_BLHD.grad is not None assert v_qat_BLHD.grad is not None - + # Check that gradients have correct shape assert q_fused_BLHD.grad.shape == q_fused_BLHD.shape assert k_fused_BLHD.grad.shape == k_fused_BLHD.shape assert v_fused_BLHD.grad.shape == v_fused_BLHD.shape - + # Check that gradients are finite assert torch.isfinite(q_fused_BLHD.grad).all() assert torch.isfinite(k_fused_BLHD.grad).all() assert torch.isfinite(v_fused_BLHD.grad).all() - + # Convert QAT gradients to float16 for comparison q_qat_grad_f16 = q_qat_BLHD.grad.to(torch.bfloat16) k_qat_grad_f16 = k_qat_BLHD.grad.to(torch.bfloat16) v_qat_grad_f16 = v_qat_BLHD.grad.to(torch.bfloat16) - + # Compare gradients dq_diff = (q_fused_BLHD.grad - q_qat_grad_f16).abs() dk_diff = (k_fused_BLHD.grad - k_qat_grad_f16).abs() dv_diff = (v_fused_BLHD.grad - v_qat_grad_f16).abs() - + dq_cos_sim = cosine_similarity(q_fused_BLHD.grad, q_qat_grad_f16) dk_cos_sim = cosine_similarity(k_fused_BLHD.grad, k_qat_grad_f16) dv_cos_sim = cosine_similarity(v_fused_BLHD.grad, v_qat_grad_f16) - - print(f" dQ - Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}") - print(f" dK - Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}") - print(f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}") - + + print( + f" dQ - Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}" + ) + print( + f" dK - Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}" + ) + print( + f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}" + ) + # Gradients may differ due to quantization in QAT and different implementations # assert dq_diff.max() < 2.0, f"dQ gradients should be reasonably close, got max_diff={dq_diff.max().item():.6f}" # assert dk_diff.max() < 2.0, f"dK gradients should be reasonably close, got max_diff={dk_diff.max().item():.6f}" # assert dv_diff.max() < 2.0, f"dV gradients should be reasonably close, got max_diff={dv_diff.max().item():.6f}" - + print("✓ Fused vs QAT backward pass test passed.") def test_fused_vs_qat_attention_different_shapes(): """Test fused vs QAT attention with different input shapes.""" torch.manual_seed(42) - + test_configs = [ (2, 4, 128, 64), # Medium - (1, 8, 256, 128), # Large head dim + (1, 8, 256, 128), # Large head dim ] - + causal = True - + for Z, H, N_CTX, HEAD_DIM in test_configs: sm_scale = 1.0 / sqrt(HEAD_DIM) - + # Create input tensors in BLHD format (float16) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) - + # Test with fused attention out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Test with QAT attention (convert to float16) q_qat_BLHD = q_BLHD.clone().to(torch.bfloat16) k_qat_BLHD = k_BLHD.clone().to(torch.bfloat16) v_qat_BLHD = v_BLHD.clone().to(torch.bfloat16) out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) out_qat_BLHD = out_qat_BLHD.to(torch.bfloat16) - + assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert torch.isfinite(out_fused_BLHD).all() assert torch.isfinite(out_qat_BLHD).all() - + # Compare outputs max_diff = (out_fused_BLHD - out_qat_BLHD).abs().max() mean_diff = (out_fused_BLHD - out_qat_BLHD).abs().mean() cos_sim = cosine_similarity(out_fused_BLHD, out_qat_BLHD) - print(f" (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # Outputs may differ due to quantization in QAT # assert max_diff < 1.0, f"Fused and QAT outputs should be reasonably close for shape (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}), got max_diff={max_diff.item():.6f}" - + print(f"✓ Fused vs QAT shape test passed for (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM})") def test_fused_vs_qat_attention_non_causal(): """Test fused vs QAT attention with non-causal masking.""" torch.manual_seed(42) - + Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = False - + # Create input tensors in BLHD format (float16) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) - + # Test with fused attention out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Test with QAT attention (convert to float16) q_qat_BLHD = q_BLHD.clone().to(torch.bfloat16) k_qat_BLHD = k_BLHD.clone().to(torch.bfloat16) v_qat_BLHD = v_BLHD.clone().to(torch.bfloat16) out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) out_qat_BLHD = out_qat_BLHD.to(torch.bfloat16) - + assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert torch.isfinite(out_fused_BLHD).all() assert torch.isfinite(out_qat_BLHD).all() - + # Compare outputs max_diff = (out_fused_BLHD - out_qat_BLHD).abs().max() mean_diff = (out_fused_BLHD - out_qat_BLHD).abs().mean() cos_sim = cosine_similarity(out_fused_BLHD, out_qat_BLHD) - print(f" Fused vs QAT (non-causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" Fused vs QAT (non-causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # Outputs may differ due to quantization in QAT # assert max_diff < 1.0, f"Fused and QAT outputs should be reasonably close for non-causal, got max_diff={max_diff.item():.6f}" - + print("✓ Fused vs QAT non-causal test passed.") def test_fused_vs_qat_attention_causal(): """Test fused vs QAT attention with causal masking.""" torch.manual_seed(42) - + Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - + # Create input tensors in BLHD format (float16) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) - + # Test with fused attention out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Test with QAT attention (convert to float16) q_qat_BLHD = q_BLHD.clone().to(torch.bfloat16) k_qat_BLHD = k_BLHD.clone().to(torch.bfloat16) v_qat_BLHD = v_BLHD.clone().to(torch.bfloat16) out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) out_qat_BLHD = out_qat_BLHD.to(torch.bfloat16) - + assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert torch.isfinite(out_fused_BLHD).all() assert torch.isfinite(out_qat_BLHD).all() - + # Compare outputs max_diff = (out_fused_BLHD - out_qat_BLHD).abs().max() mean_diff = (out_fused_BLHD - out_qat_BLHD).abs().mean() cos_sim = cosine_similarity(out_fused_BLHD, out_qat_BLHD) - print(f" Fused vs QAT (causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" Fused vs QAT (causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # Outputs may differ due to quantization in QAT # assert max_diff < 1.0, f"Fused and QAT outputs should be reasonably close for causal, got max_diff={max_diff.item():.6f}" - + print("✓ Fused vs QAT causal test passed.") def test_fused_vs_qat_attention_causal_backward(): """Test backward pass of fused vs QAT attention with causal masking.""" torch.manual_seed(42) - + Z, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - + # Create input tensors in BLHD format for fused attention (float16) q_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True) k_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True) v_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True) - + # Create input tensors for QAT (float16, same values) q_qat_BLHD = q_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True) k_qat_BLHD = k_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True) v_qat_BLHD = v_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True) - + # Forward pass with fused attention out_fused_BLHD = fused_attn_wrapper(q_fused_BLHD, k_fused_BLHD, v_fused_BLHD, causal, sm_scale) - + # Forward pass with QAT out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) - + # Create dummy gradient dout = torch.randn_like(out_fused_BLHD).contiguous() dout_qat = dout.to(torch.bfloat16) - + # Backward pass for fused attention out_fused_BLHD.backward(dout) - + # Backward pass for QAT out_qat_BLHD.backward(dout_qat) - + # Check that gradients are computed assert q_fused_BLHD.grad is not None assert k_fused_BLHD.grad is not None @@ -1194,48 +1257,54 @@ def test_fused_vs_qat_attention_causal_backward(): assert q_qat_BLHD.grad is not None assert k_qat_BLHD.grad is not None assert v_qat_BLHD.grad is not None - + # Convert QAT gradients to float16 for comparison q_qat_grad_f16 = q_qat_BLHD.grad.to(torch.bfloat16) k_qat_grad_f16 = k_qat_BLHD.grad.to(torch.bfloat16) v_qat_grad_f16 = v_qat_BLHD.grad.to(torch.bfloat16) - + # Compare gradients dq_diff = (q_fused_BLHD.grad - q_qat_grad_f16).abs() dk_diff = (k_fused_BLHD.grad - k_qat_grad_f16).abs() dv_diff = (v_fused_BLHD.grad - v_qat_grad_f16).abs() - + dq_cos_sim = cosine_similarity(q_fused_BLHD.grad, q_qat_grad_f16) dk_cos_sim = cosine_similarity(k_fused_BLHD.grad, k_qat_grad_f16) dv_cos_sim = cosine_similarity(v_fused_BLHD.grad, v_qat_grad_f16) - - print(f" dQ - Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}") - print(f" dK - Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}") - print(f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}") - + + print( + f" dQ - Max diff: {dq_diff.max().item():.6f}, Mean diff: {dq_diff.mean().item():.6f}, Cosine sim: {dq_cos_sim:.6f}" + ) + print( + f" dK - Max diff: {dk_diff.max().item():.6f}, Mean diff: {dk_diff.mean().item():.6f}, Cosine sim: {dk_cos_sim:.6f}" + ) + print( + f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}" + ) + print("✓ Fused vs QAT causal backward test passed.") def test_qat_attention_wan_shape_forward(): """Test QAT attention forward pass with WAN shape [1, 40, 9360, 128].""" torch.manual_seed(42) - + # WAN shape: [1, 40, 9360, 128] = (Z, H, N_CTX, HEAD_DIM) Z, H, N_CTX, HEAD_DIM = 1, 40, 9360, 128 dtype = torch.bfloat16 sm_scale = 1.0 / sqrt(HEAD_DIM) causal = False # WAN uses non-causal attention - + print(f" Testing WAN shape: Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}") - + # Create input tensors in BLHD format (B, L, H, D) = (Z, N_CTX, H, HEAD_DIM) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) - + # Test with QAT (using wrapper that does permute/contiguous) out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - + # Convert to BHLD for naive attention comparison q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() @@ -1246,21 +1315,23 @@ def test_qat_attention_wan_shape_forward(): print(f" out_naive_BLHD shape (after permute): {out_naive_BLHD.shape}") print(f" Expected shape: {(Z, N_CTX, H, HEAD_DIM)}") print(f" out_qat_BLHD shape: {out_qat_BLHD.shape}") - + # Check that outputs have correct shape assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) - + # Check that outputs are finite assert torch.isfinite(out_qat_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() - + # Compare QAT and naive outputs max_diff = (out_qat_BLHD - out_naive_BLHD).abs().max() mean_diff = (out_qat_BLHD - out_naive_BLHD).abs().mean() cos_sim = cosine_similarity(out_qat_BLHD, out_naive_BLHD) - print(f" QAT vs Naive (WAN shape) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - + print( + f" QAT vs Naive (WAN shape) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" + ) + # QAT output should be reasonably close to naive (within tolerance due to quantization) print("✓ WAN shape forward pass test passed.") @@ -1286,9 +1357,9 @@ def test_qat_attention_wan_shape_backward(): # ----------------------------- # CREATE BHLD INPUTS FOR NAIVE # ----------------------------- - q_BHLD = q_BLHD.detach().permute(0,2,1,3).contiguous().requires_grad_(True) - k_BHLD = k_BLHD.detach().permute(0,2,1,3).contiguous().requires_grad_(True) - v_BHLD = v_BLHD.detach().permute(0,2,1,3).contiguous().requires_grad_(True) + q_BHLD = q_BLHD.detach().permute(0, 2, 1, 3).contiguous().requires_grad_(True) + k_BHLD = k_BLHD.detach().permute(0, 2, 1, 3).contiguous().requires_grad_(True) + v_BHLD = v_BLHD.detach().permute(0, 2, 1, 3).contiguous().requires_grad_(True) # ----------------------------- # FORWARD — QAT @@ -1301,7 +1372,7 @@ def test_qat_attention_wan_shape_backward(): out_naive_BHLD = naive_attention(q_BHLD, k_BHLD, v_BHLD, causal, sm_scale) # Convert naive output back to BLHD for comparing outputs - out_naive_BLHD = out_naive_BHLD.permute(0,2,1,3).contiguous() + out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() # ----------------------------- # BACKWARD @@ -1312,7 +1383,7 @@ def test_qat_attention_wan_shape_backward(): out_qat_BLHD.backward(dout_BLHD) # ❗ FIXED: convert BLHD dout → BHLD for naive backward - dout_BHLD = dout_BLHD.permute(0,2,1,3).contiguous() + dout_BHLD = dout_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD.backward(dout_BHLD) @@ -1320,23 +1391,35 @@ def test_qat_attention_wan_shape_backward(): # GRADIENT FORMAT FIX # Convert naive BHLD grads → BLHD # ----------------------------- - q_naive_grad_BLHD = q_BHLD.grad.permute(0,2,1,3).contiguous() - k_naive_grad_BLHD = k_BHLD.grad.permute(0,2,1,3).contiguous() - v_naive_grad_BLHD = v_BHLD.grad.permute(0,2,1,3).contiguous() + q_naive_grad_BLHD = q_BHLD.grad.permute(0, 2, 1, 3).contiguous() + k_naive_grad_BLHD = k_BHLD.grad.permute(0, 2, 1, 3).contiguous() + v_naive_grad_BLHD = v_BHLD.grad.permute(0, 2, 1, 3).contiguous() # ----------------------------- # PRINT GRADIENTS # ----------------------------- print("\n === QAT Gradients ===") - print(f" dQ_QAT - Shape: {q_BLHD.grad.shape}, Min: {q_BLHD.grad.min().item():.6f}, Max: {q_BLHD.grad.max().item():.6f}, Mean: {q_BLHD.grad.mean().item():.6f}, Std: {q_BLHD.grad.std().item():.6f}") - print(f" dK_QAT - Shape: {k_BLHD.grad.shape}, Min: {k_BLHD.grad.min().item():.6f}, Max: {k_BLHD.grad.max().item():.6f}, Mean: {k_BLHD.grad.mean().item():.6f}, Std: {k_BLHD.grad.std().item():.6f}") - print(f" dV_QAT - Shape: {v_BLHD.grad.shape}, Min: {v_BLHD.grad.min().item():.6f}, Max: {v_BLHD.grad.max().item():.6f}, Mean: {v_BLHD.grad.mean().item():.6f}, Std: {v_BLHD.grad.std().item():.6f}") - + print( + f" dQ_QAT - Shape: {q_BLHD.grad.shape}, Min: {q_BLHD.grad.min().item():.6f}, Max: {q_BLHD.grad.max().item():.6f}, Mean: {q_BLHD.grad.mean().item():.6f}, Std: {q_BLHD.grad.std().item():.6f}" + ) + print( + f" dK_QAT - Shape: {k_BLHD.grad.shape}, Min: {k_BLHD.grad.min().item():.6f}, Max: {k_BLHD.grad.max().item():.6f}, Mean: {k_BLHD.grad.mean().item():.6f}, Std: {k_BLHD.grad.std().item():.6f}" + ) + print( + f" dV_QAT - Shape: {v_BLHD.grad.shape}, Min: {v_BLHD.grad.min().item():.6f}, Max: {v_BLHD.grad.max().item():.6f}, Mean: {v_BLHD.grad.mean().item():.6f}, Std: {v_BLHD.grad.std().item():.6f}" + ) + print("\n === Naive Gradients ===") - print(f" dQ_Naive - Shape: {q_naive_grad_BLHD.shape}, Min: {q_naive_grad_BLHD.min().item():.6f}, Max: {q_naive_grad_BLHD.max().item():.6f}, Mean: {q_naive_grad_BLHD.mean().item():.6f}, Std: {q_naive_grad_BLHD.std().item():.6f}") - print(f" dK_Naive - Shape: {k_naive_grad_BLHD.shape}, Min: {k_naive_grad_BLHD.min().item():.6f}, Max: {k_naive_grad_BLHD.max().item():.6f}, Mean: {k_naive_grad_BLHD.mean().item():.6f}, Std: {k_naive_grad_BLHD.std().item():.6f}") - print(f" dV_Naive - Shape: {v_naive_grad_BLHD.shape}, Min: {v_naive_grad_BLHD.min().item():.6f}, Max: {v_naive_grad_BLHD.max().item():.6f}, Mean: {v_naive_grad_BLHD.mean().item():.6f}, Std: {v_naive_grad_BLHD.std().item():.6f}") - + print( + f" dQ_Naive - Shape: {q_naive_grad_BLHD.shape}, Min: {q_naive_grad_BLHD.min().item():.6f}, Max: {q_naive_grad_BLHD.max().item():.6f}, Mean: {q_naive_grad_BLHD.mean().item():.6f}, Std: {q_naive_grad_BLHD.std().item():.6f}" + ) + print( + f" dK_Naive - Shape: {k_naive_grad_BLHD.shape}, Min: {k_naive_grad_BLHD.min().item():.6f}, Max: {k_naive_grad_BLHD.max().item():.6f}, Mean: {k_naive_grad_BLHD.mean().item():.6f}, Std: {k_naive_grad_BLHD.std().item():.6f}" + ) + print( + f" dV_Naive - Shape: {v_naive_grad_BLHD.shape}, Min: {v_naive_grad_BLHD.min().item():.6f}, Max: {v_naive_grad_BLHD.max().item():.6f}, Mean: {v_naive_grad_BLHD.mean().item():.6f}, Std: {v_naive_grad_BLHD.std().item():.6f}" + ) + # Print sample values (first few elements) # print("\n === Sample Gradient Values (first 10 elements) ===") # print(f" dQ_QAT[0,0,0,:10]: {q_BLHD.grad[0,0,0,:10].cpu().tolist()}") @@ -1365,7 +1448,6 @@ def test_qat_attention_wan_shape_backward(): print("\n✓ WAN shape backward pass test passed.") - def test_qat_attention_wan_shape_non_divisible_64(): """Test QAT attention where N_CTX is not divisible by 64.""" torch.manual_seed(42) @@ -1388,16 +1470,16 @@ def test_qat_attention_wan_shape_non_divisible_64(): # ----------------------------- # BHLD for naive # ----------------------------- - q_BHLD = q_BLHD.detach().permute(0,2,1,3).contiguous().requires_grad_(True) - k_BHLD = k_BLHD.detach().permute(0,2,1,3).contiguous().requires_grad_(True) - v_BHLD = v_BLHD.detach().permute(0,2,1,3).contiguous().requires_grad_(True) + q_BHLD = q_BLHD.detach().permute(0, 2, 1, 3).contiguous().requires_grad_(True) + k_BHLD = k_BLHD.detach().permute(0, 2, 1, 3).contiguous().requires_grad_(True) + v_BHLD = v_BLHD.detach().permute(0, 2, 1, 3).contiguous().requires_grad_(True) # ----------------------------- # FORWARD # ----------------------------- out_qat_BLHD = attn_qat_train_wrapper(q_BLHD, k_BLHD, v_BLHD, causal, sm_scale) out_naive_BHLD = naive_attention(q_BHLD, k_BHLD, v_BHLD, causal, sm_scale) - out_naive_BLHD = out_naive_BHLD.permute(0,2,1,3).contiguous() + out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() # ----------------------------- # BACKWARD @@ -1408,30 +1490,42 @@ def test_qat_attention_wan_shape_non_divisible_64(): out_qat_BLHD.backward(dout_BLHD) # FIXED: Convert dout to BHLD for naive - dout_BHLD = dout_BLHD.permute(0,2,1,3).contiguous() + dout_BHLD = dout_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD.backward(dout_BHLD) # ----------------------------- # GRADIENT FORMAT FIX # ----------------------------- - q_naive_grad_BLHD = q_BHLD.grad.permute(0,2,1,3).contiguous() - k_naive_grad_BLHD = k_BHLD.grad.permute(0,2,1,3).contiguous() - v_naive_grad_BLHD = v_BHLD.grad.permute(0,2,1,3).contiguous() + q_naive_grad_BLHD = q_BHLD.grad.permute(0, 2, 1, 3).contiguous() + k_naive_grad_BLHD = k_BHLD.grad.permute(0, 2, 1, 3).contiguous() + v_naive_grad_BLHD = v_BHLD.grad.permute(0, 2, 1, 3).contiguous() # ----------------------------- # PRINT GRADIENTS # ----------------------------- print("\n === QAT Gradients ===") - print(f" dQ_QAT - Shape: {q_BLHD.grad.shape}, Min: {q_BLHD.grad.min().item():.6f}, Max: {q_BLHD.grad.max().item():.6f}, Mean: {q_BLHD.grad.mean().item():.6f}, Std: {q_BLHD.grad.std().item():.6f}") - print(f" dK_QAT - Shape: {k_BLHD.grad.shape}, Min: {k_BLHD.grad.min().item():.6f}, Max: {k_BLHD.grad.max().item():.6f}, Mean: {k_BLHD.grad.mean().item():.6f}, Std: {k_BLHD.grad.std().item():.6f}") - print(f" dV_QAT - Shape: {v_BLHD.grad.shape}, Min: {v_BLHD.grad.min().item():.6f}, Max: {v_BLHD.grad.max().item():.6f}, Mean: {v_BLHD.grad.mean().item():.6f}, Std: {v_BLHD.grad.std().item():.6f}") - + print( + f" dQ_QAT - Shape: {q_BLHD.grad.shape}, Min: {q_BLHD.grad.min().item():.6f}, Max: {q_BLHD.grad.max().item():.6f}, Mean: {q_BLHD.grad.mean().item():.6f}, Std: {q_BLHD.grad.std().item():.6f}" + ) + print( + f" dK_QAT - Shape: {k_BLHD.grad.shape}, Min: {k_BLHD.grad.min().item():.6f}, Max: {k_BLHD.grad.max().item():.6f}, Mean: {k_BLHD.grad.mean().item():.6f}, Std: {k_BLHD.grad.std().item():.6f}" + ) + print( + f" dV_QAT - Shape: {v_BLHD.grad.shape}, Min: {v_BLHD.grad.min().item():.6f}, Max: {v_BLHD.grad.max().item():.6f}, Mean: {v_BLHD.grad.mean().item():.6f}, Std: {v_BLHD.grad.std().item():.6f}" + ) + print("\n === Naive Gradients ===") - print(f" dQ_Naive - Shape: {q_naive_grad_BLHD.shape}, Min: {q_naive_grad_BLHD.min().item():.6f}, Max: {q_naive_grad_BLHD.max().item():.6f}, Mean: {q_naive_grad_BLHD.mean().item():.6f}, Std: {q_naive_grad_BLHD.std().item():.6f}") - print(f" dK_Naive - Shape: {k_naive_grad_BLHD.shape}, Min: {k_naive_grad_BLHD.min().item():.6f}, Max: {k_naive_grad_BLHD.max().item():.6f}, Mean: {k_naive_grad_BLHD.mean().item():.6f}, Std: {k_naive_grad_BLHD.std().item():.6f}") - print(f" dV_Naive - Shape: {v_naive_grad_BLHD.shape}, Min: {v_naive_grad_BLHD.min().item():.6f}, Max: {v_naive_grad_BLHD.max().item():.6f}, Mean: {v_naive_grad_BLHD.mean().item():.6f}, Std: {v_naive_grad_BLHD.std().item():.6f}") - + print( + f" dQ_Naive - Shape: {q_naive_grad_BLHD.shape}, Min: {q_naive_grad_BLHD.min().item():.6f}, Max: {q_naive_grad_BLHD.max().item():.6f}, Mean: {q_naive_grad_BLHD.mean().item():.6f}, Std: {q_naive_grad_BLHD.std().item():.6f}" + ) + print( + f" dK_Naive - Shape: {k_naive_grad_BLHD.shape}, Min: {k_naive_grad_BLHD.min().item():.6f}, Max: {k_naive_grad_BLHD.max().item():.6f}, Mean: {k_naive_grad_BLHD.mean().item():.6f}, Std: {k_naive_grad_BLHD.std().item():.6f}" + ) + print( + f" dV_Naive - Shape: {v_naive_grad_BLHD.shape}, Min: {v_naive_grad_BLHD.min().item():.6f}, Max: {v_naive_grad_BLHD.max().item():.6f}, Mean: {v_naive_grad_BLHD.mean().item():.6f}, Std: {v_naive_grad_BLHD.std().item():.6f}" + ) + # Print sample values (first few elements) # print("\n === Sample Gradient Values (first 10 elements) ===") # print(f" dQ_QAT[0,0,0,:10]: {q_BLHD.grad[0,0,0,:10].cpu().tolist()}") @@ -1464,7 +1558,7 @@ if __name__ == "__main__": print("Running QAT attention tests...") print(f"Device: {DEVICE}") print() - + test_qat_attention_forward() test_qat_attention_backward() test_qat_attention_different_shapes() @@ -1476,28 +1570,28 @@ if __name__ == "__main__": test_qat_attention_wan_shape_forward() test_qat_attention_wan_shape_backward() test_qat_attention_wan_shape_non_divisible_64() - + print() print("Running Fused attention tests...") print() - + # test_fused_attention_forward() # test_fused_attention_backward() # test_fused_attention_different_shapes() # test_fused_attention_non_causal() # test_fused_attention_causal() # test_fused_attention_causal_backward() - + # print() # print("Running Fused vs QAT attention comparison tests...") # print() - + # test_fused_vs_qat_attention_forward() # test_fused_vs_qat_attention_backward() # test_fused_vs_qat_attention_different_shapes() # test_fused_vs_qat_attention_non_causal() # test_fused_vs_qat_attention_causal() # test_fused_vs_qat_attention_causal_backward() - + print() print("All tests passed! ✓") diff --git a/fastvideo-kernel/tests/test_fused_compress_topk.py b/fastvideo-kernel/tests/test_fused_compress_topk.py index db0afd40a..adbbbd632 100644 --- a/fastvideo-kernel/tests/test_fused_compress_topk.py +++ b/fastvideo-kernel/tests/test_fused_compress_topk.py @@ -34,9 +34,7 @@ class TestFusedTopkMaskTies: for topk in [1, 3, 5, 8]: mask = self._run_topk_mask(scores, topk) row_counts = mask.sum(dim=-1) - assert (row_counts == topk).all(), ( - f"topk={topk}, row_counts={row_counts}" - ) + assert (row_counts == topk).all(), (f"topk={topk}, row_counts={row_counts}") def test_reviewer_example(self): """Exact case from the reviewer: two 0.8875 with topk=1 → must select 1.""" @@ -58,9 +56,7 @@ class TestFusedTopkMaskTies: ).view(1, 1, 1, -1) for topk in [1, 2, 3, 4, 7, 8]: mask = self._run_topk_mask(scores, topk) - assert mask.sum().item() == topk, ( - f"topk={topk}, selected={mask.sum().item()}" - ) + assert mask.sum().item() == topk, (f"topk={topk}, selected={mask.sum().item()}") def test_batch_heads_ties(self): """Multiple batch/heads with tie-heavy scores.""" @@ -73,8 +69,7 @@ class TestFusedTopkMaskTies: mask = self._run_topk_mask(scores, topk) row_counts = mask.sum(dim=-1) assert (row_counts == topk).all(), ( - f"topk={topk}, row_counts min={row_counts.min()}, max={row_counts.max()}" - ) + f"topk={topk}, row_counts min={row_counts.min()}, max={row_counts.max()}") # --------------------------------------------------------------------------- @@ -132,7 +127,7 @@ class TestFusedBlockMeanEquivalence: seq_len = num_blocks * block_elements x = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device="cuda") - vbs = torch.randint(16, 65, (num_blocks,), dtype=torch.int32, device="cuda") + vbs = torch.randint(16, 65, (num_blocks, ), dtype=torch.int32, device="cuda") fused_out = fused_block_mean(x, vbs, block_elements) @@ -140,9 +135,8 @@ class TestFusedBlockMeanEquivalence: ref_out = (x_blocks.float().sum(dim=3) / vbs.view(1, 1, -1, 1).float()).to(torch.bfloat16) assert fused_out.shape == ref_out.shape - assert torch.allclose(fused_out, ref_out, atol=1e-2, rtol=1e-2), ( - f"max diff={( fused_out - ref_out).abs().max().item()}" - ) + assert torch.allclose(fused_out, ref_out, atol=1e-2, + rtol=1e-2), (f"max diff={( fused_out - ref_out).abs().max().item()}") # --------------------------------------------------------------------------- @@ -170,7 +164,7 @@ class TestFusedBlockMeanBackward: block_elements = 64 seq_len = num_blocks * block_elements - vbs = torch.randint(16, 65, (num_blocks,), dtype=torch.int32, device="cuda") + vbs = torch.randint(16, 65, (num_blocks, ), dtype=torch.int32, device="cuda") grad_out = torch.randn(B, H, num_blocks, D, dtype=torch.bfloat16, device="cuda") x_fused = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device="cuda") @@ -189,9 +183,7 @@ class TestFusedBlockMeanBackward: assert x_fused.grad.shape == x_ref.grad.shape max_diff = (x_fused.grad - x_ref.grad).abs().max().item() - assert torch.allclose(x_fused.grad, x_ref.grad, atol=1e-2, rtol=1e-2), ( - f"gradient max diff={max_diff}" - ) + assert torch.allclose(x_fused.grad, x_ref.grad, atol=1e-2, rtol=1e-2), (f"gradient max diff={max_diff}") def test_gradient_parity_variable_block_sizes(self): """Backward parity with non-uniform variable_block_sizes.""" @@ -203,7 +195,7 @@ class TestFusedBlockMeanBackward: block_elements = 64 seq_len = num_blocks * block_elements - vbs = torch.randint(1, 65, (num_blocks,), dtype=torch.int32, device="cuda") + vbs = torch.randint(1, 65, (num_blocks, ), dtype=torch.int32, device="cuda") grad_out = torch.randn(B, H, num_blocks, D, dtype=torch.bfloat16, device="cuda") x_fused = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device="cuda") @@ -218,9 +210,7 @@ class TestFusedBlockMeanBackward: out_ref.backward(grad_out) max_diff = (x_fused.grad - x_ref.grad).abs().max().item() - assert torch.allclose(x_fused.grad, x_ref.grad, atol=1e-2, rtol=1e-2), ( - f"gradient max diff={max_diff}" - ) + assert torch.allclose(x_fused.grad, x_ref.grad, atol=1e-2, rtol=1e-2), (f"gradient max diff={max_diff}") def test_gradient_through_matmul_chain(self): """End-to-end backward through the compress branch: block_mean → matmul → softmax → matmul.""" @@ -232,16 +222,14 @@ class TestFusedBlockMeanBackward: block_elements = 64 seq_len = num_blocks * block_elements - vbs = torch.full((num_blocks,), block_elements, dtype=torch.int32, device="cuda") + vbs = torch.full((num_blocks, ), block_elements, dtype=torch.int32, device="cuda") def _forward_compress(x_q, x_k, x_v, use_fused): - mean_fn = fused_block_mean if use_fused else ( - lambda x, v, be: _ref_block_mean_with_grad(x, v, be) - ) + mean_fn = fused_block_mean if use_fused else (lambda x, v, be: _ref_block_mean_with_grad(x, v, be)) q_c = mean_fn(x_q, vbs, block_elements) k_c = mean_fn(x_k, vbs, block_elements) v_c = mean_fn(x_v, vbs, block_elements) - scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (D ** 0.5) + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (D**0.5) attn = torch.softmax(scores, dim=-1) return torch.matmul(attn, v_c) @@ -261,12 +249,10 @@ class TestFusedBlockMeanBackward: for name, (g_fused, g_ref) in zip( ["dQ", "dK", "dV"], - zip(grads["fused"], grads["ref"]), + zip(grads["fused"], grads["ref"]), ): max_diff = (g_fused - g_ref).abs().max().item() - assert torch.allclose(g_fused, g_ref, atol=1e-2, rtol=1e-2), ( - f"{name} gradient max diff={max_diff}" - ) + assert torch.allclose(g_fused, g_ref, atol=1e-2, rtol=1e-2), (f"{name} gradient max diff={max_diff}") # --------------------------------------------------------------------------- @@ -275,8 +261,13 @@ class TestFusedBlockMeanBackward: def _old_pytorch_video_sparse_attn( - q, k, v, variable_block_sizes, q_variable_block_sizes, - topk, block_elements, + q, + k, + v, + variable_block_sizes, + q_variable_block_sizes, + topk, + block_elements, ): """Reproduce the pre-optimization PyTorch pipeline from ops.py.""" batch, heads, q_seq_len, dim = q.shape @@ -291,7 +282,7 @@ def _old_pytorch_video_sparse_attn( k_c = (k_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(k.dtype) v_c = (v_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(v.dtype) - scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim ** 0.5) + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim**0.5) attn = torch.softmax(scores, dim=-1) out_c = torch.matmul(attn, v_c) out_c = out_c.view(batch, heads, q_num_blocks, 1, dim) @@ -311,7 +302,7 @@ class TestVideoSparseAttnEquivalence: q = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device=device) k = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device=device) v = torch.randn(B, H, seq_len, D, dtype=torch.bfloat16, device=device) - vbs = torch.full((num_blocks,), block_elements, dtype=torch.int32, device=device) + vbs = torch.full((num_blocks, ), block_elements, dtype=torch.int32, device=device) return q, k, v, vbs def test_compress_branch_equivalence(self): @@ -327,17 +318,22 @@ class TestVideoSparseAttnEquivalence: q_c_fused = fused_block_mean(q, vbs, block_elements) k_c_fused = fused_block_mean(k, vbs, block_elements) v_c_fused = fused_block_mean(v, vbs, block_elements) - scores_fused = torch.matmul(q_c_fused, k_c_fused.transpose(-2, -1)) / (D ** 0.5) + scores_fused = torch.matmul(q_c_fused, k_c_fused.transpose(-2, -1)) / (D**0.5) attn_fused = torch.softmax(scores_fused, dim=-1) out_c_fused = torch.matmul(attn_fused, v_c_fused) old_out_c, old_scores, _ = _old_pytorch_video_sparse_attn( - q, k, v, vbs, vbs, topk, block_elements, + q, + k, + v, + vbs, + vbs, + topk, + block_elements, ) - assert torch.allclose(scores_fused, old_scores, atol=1e-2, rtol=1e-2), ( - f"scores max diff={( scores_fused - old_scores).abs().max().item()}" - ) + assert torch.allclose(scores_fused, old_scores, atol=1e-2, + rtol=1e-2), (f"scores max diff={( scores_fused - old_scores).abs().max().item()}") def test_topk_mask_row_count_matches(self): """Fused topk mask must select exactly topk per row, same as torch.topk.""" @@ -354,16 +350,20 @@ class TestVideoSparseAttnEquivalence: q_c = fused_block_mean(q, vbs, block_elements) k_c = fused_block_mean(k, vbs, block_elements) - scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (D ** 0.5) + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (D**0.5) fused_mask = fused_topk_mask(scores, topk) row_counts = fused_mask.sum(dim=-1) - assert (row_counts == topk).all(), ( - f"row_counts min={row_counts.min()}, max={row_counts.max()}" - ) + assert (row_counts == topk).all(), (f"row_counts min={row_counts.min()}, max={row_counts.max()}") _, _, ref_mask = _old_pytorch_video_sparse_attn( - q, k, v, vbs, vbs, topk, block_elements, + q, + k, + v, + vbs, + vbs, + topk, + block_elements, ) ref_counts = ref_mask.sum(dim=-1) assert (ref_counts == topk).all() @@ -383,7 +383,7 @@ class TestVideoSparseAttnEquivalence: q_c = fused_block_mean(q, vbs, block_elements) k_c = fused_block_mean(k, vbs, block_elements) - scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (D ** 0.5) + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (D**0.5) fused_mask = fused_topk_mask(scores, topk) diff --git a/fastvideo-kernel/tests/test_sta.py b/fastvideo-kernel/tests/test_sta.py index 1402083f2..02b0bd4ef 100644 --- a/fastvideo-kernel/tests/test_sta.py +++ b/fastvideo-kernel/tests/test_sta.py @@ -85,11 +85,13 @@ def test_sliding_tile_attention(): # Run correctness check directly results = check_correctness(b, h, n, d, causal, mean, std, error_mode='output') - assert results['TK vs FLEX']['avg_diff'] < 3e-6, f"Average difference: {results['TK vs FLEX']['avg_diff']} is too large" - assert results['TK vs FLEX']['max_diff'] < 4e-2, f"Maximum difference: {results['TK vs FLEX']['max_diff']} is too large" + assert results['TK vs FLEX'][ + 'avg_diff'] < 3e-6, f"Average difference: {results['TK vs FLEX']['avg_diff']} is too large" + assert results['TK vs FLEX'][ + 'max_diff'] < 4e-2, f"Maximum difference: {results['TK vs FLEX']['max_diff']} is too large" print(f"Average difference: {results['TK vs FLEX']['avg_diff']}") print(f"Maximum difference: {results['TK vs FLEX']['max_diff']}") + if __name__ == "__main__": test_sliding_tile_attention() - diff --git a/fastvideo-kernel/tests/test_turbodiffusion.py b/fastvideo-kernel/tests/test_turbodiffusion.py index 2233097ff..771ab7ffa 100644 --- a/fastvideo-kernel/tests/test_turbodiffusion.py +++ b/fastvideo-kernel/tests/test_turbodiffusion.py @@ -13,6 +13,7 @@ except ImportError: from fastvideo_kernel import turbodiffusion_ops + # Helper for RMS Norm reference def rms_norm_ref(x, w, eps=1e-6): dtype = x.dtype @@ -21,56 +22,57 @@ def rms_norm_ref(x, w, eps=1e-6): x = x * torch.rsqrt(variance + eps) return (x * w.float()).to(dtype) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") class TestTurboDiffusion: - + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("shape", [(16, 128), (32, 256), (1, 1024)]) def test_quant_correctness(self, dtype, shape): if turbodiffusion_ops.quant_cuda is None: pytest.skip("quant_cuda not available") - + x = torch.randn(shape, dtype=dtype, device="cuda") x_q, x_scale = turbodiffusion_ops.int8_quant(x) - + assert x_q.dtype == torch.int8 assert x_scale.dtype == torch.float32 - + # Simple check: dequantize and compute error # Note: The quantization scheme details matter here (per block? per tensor?). - # Looking at quant.cu, it seems to be block-based but the output scale shape isn't immediately obvious from python signature + # Looking at quant.cu, it seems to be block-based but the output scale shape isn't immediately obvious from python signature # without looking at C++ code deeper. # But let's check shapes at least. - - # If we can't easily dequantize without knowing block size logic in python, + + # If we can't easily dequantize without knowing block size logic in python, # checking that it runs and produces valid shapes is a good start. assert x_q.shape == shape # x_scale shape depends on block size, usually smaller than x - + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_gemm_correctness(self, dtype): if turbodiffusion_ops.gemm_cuda is None: pytest.skip("gemm_cuda not available") - + M, N, K = 32, 64, 128 x = torch.randn(M, K, dtype=dtype, device="cuda") - + # Create weights # For simplicity in testing, let's create random int8 weights and scales w_q = torch.randint(-127, 127, (N, K), dtype=torch.int8, device="cuda") - + # Scale shape: The Int8Linear class uses: # row_blocks = cdiv(out_features, b=128) # col_blocks = cdiv(in_features, b=128) # scale shape: (row_blocks, col_blocks) - + row_blocks = (N + 127) // 128 col_blocks = (K + 127) // 128 w_s = torch.randn(row_blocks, col_blocks, dtype=torch.float32, device="cuda").abs() - + # Run int8_linear output = turbodiffusion_ops.int8_linear(x, w_q, w_s) - + assert output.shape == (M, N) assert output.dtype == dtype @@ -78,48 +80,47 @@ class TestTurboDiffusion: def test_gemm_backward(self, dtype): if turbodiffusion_ops.gemm_cuda is None: pytest.skip("gemm_cuda not available") - + M, N, K = 32, 64, 128 x = torch.randn(M, K, dtype=dtype, device="cuda", requires_grad=True) - + # Weights (frozen) w_q = torch.randint(-127, 127, (N, K), dtype=torch.int8, device="cuda") row_blocks = (N + 127) // 128 col_blocks = (K + 127) // 128 w_s = torch.randn(row_blocks, col_blocks, dtype=torch.float32, device="cuda").abs() - + bias = torch.randn(N, dtype=dtype, device="cuda", requires_grad=True) - + # Use Int8LinearFunction output = turbodiffusion_ops.Int8LinearFunction.apply(x, w_q, w_s, bias) loss = output.sum() loss.backward() - + assert x.grad is not None assert bias.grad is not None assert x.grad.shape == (M, K) - + # Check correctness against dequantized weight w_float = turbodiffusion_ops.dequantize_weight(w_q, w_s, dtype) - + # Reference x_ref = x.detach().clone().requires_grad_() bias_ref = bias.detach().clone().requires_grad_() - + # Note: Int8Linear forward uses quantized input, so output won't match exactly reference with float input. # But backward gradients should be consistent with the logic we implemented: # grad_input = grad_output @ w_float - + # Let's verify our backward logic matches standard matmul backward with dequantized weights # We manually compute expected gradient given grad_output = ones grad_output = torch.ones_like(output) expected_x_grad = grad_output @ w_float expected_bias_grad = grad_output.sum(0) - + torch.testing.assert_close(x.grad, expected_x_grad, atol=1e-2, rtol=1e-2) torch.testing.assert_close(bias.grad, expected_bias_grad, atol=1e-2, rtol=1e-2) - @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("shape", [(2, 16, 128), (4, 32, 256)]) def test_rms_norm_triton(self, dtype, shape): @@ -127,7 +128,7 @@ class TestTurboDiffusion: dim = shape[-1] w = torch.randn(dim, dtype=dtype, device="cuda") eps = 1e-5 - + # Triton implementation # Note: rmsnorm returns tuple now res = turbodiffusion_ops.rmsnorm(x, w, eps) @@ -135,10 +136,10 @@ class TestTurboDiffusion: out_triton = res[0] else: out_triton = res - + # Reference out_ref = rms_norm_ref(x, w, eps) - + torch.testing.assert_close(out_triton, out_ref, atol=1e-2, rtol=1e-2) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -148,29 +149,30 @@ class TestTurboDiffusion: dim = shape[-1] w = torch.randn(dim, dtype=dtype, device="cuda", requires_grad=True) eps = 1e-5 - + # Forward via Function y = turbodiffusion_ops.RMSNormFunction.apply(x, w, eps) loss = y.sum() loss.backward() - + x_grad = x.grad w_grad = w.grad - + # Reference x_ref = x.detach().clone().requires_grad_() w_ref = w.detach().clone().requires_grad_() + # Custom RMSNorm ref in pytorch def rms_norm_ref_grad(x, w, eps): - x_float = x.float() - var = x_float.pow(2).mean(-1, keepdim=True) - rstd = torch.rsqrt(var + eps) - return (x_float * rstd * w.float()).to(x.dtype) - + x_float = x.float() + var = x_float.pow(2).mean(-1, keepdim=True) + rstd = torch.rsqrt(var + eps) + return (x_float * rstd * w.float()).to(x.dtype) + y_ref = rms_norm_ref_grad(x_ref, w_ref, eps) loss_ref = y_ref.sum() loss_ref.backward() - + torch.testing.assert_close(x_grad, x_ref.grad, atol=1e-2, rtol=1e-2) torch.testing.assert_close(w_grad, w_ref.grad, atol=2e-2, rtol=2e-2) @@ -178,21 +180,21 @@ class TestTurboDiffusion: def test_rms_norm_cuda(self, dtype): if fastvideo_kernel_ops is None or not hasattr(fastvideo_kernel_ops, "rms_norm_cuda"): pytest.skip("rms_norm_cuda not available") - + shape = (16, 128) x = torch.randn(shape, dtype=dtype, device="cuda") dim = shape[-1] w = torch.randn(dim, dtype=dtype, device="cuda") eps = 1e-5 - + # C++ implementation # Signature: rms_norm_cuda(Input, eps, Weight, Output) -> Output out_cuda = torch.empty_like(x) fastvideo_kernel_ops.rms_norm_cuda(x, eps, w, out_cuda) - + # Reference out_ref = rms_norm_ref(x, w, eps) - + torch.testing.assert_close(out_cuda, out_ref, atol=1e-2, rtol=1e-2) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -201,11 +203,11 @@ class TestTurboDiffusion: x = torch.randn(shape, dtype=dtype, device="cuda") dim = shape[-1] eps = 1e-5 - + # With affine w = torch.randn(dim, dtype=dtype, device="cuda") b = torch.randn(dim, dtype=dtype, device="cuda") - + # Triton implementation # Note: layernorm returns tuple now res = turbodiffusion_ops.layernorm(x, w, b, eps, elementwise_affine=True) @@ -214,15 +216,15 @@ class TestTurboDiffusion: else: out_triton = res out_triton = out_triton.to(dtype) - + # Reference ln = torch.nn.LayerNorm(dim, eps=eps, elementwise_affine=True, dtype=dtype).cuda() ln.weight.data.copy_(w) ln.bias.data.copy_(b) out_ref = ln(x) - + torch.testing.assert_close(out_triton, out_ref, atol=1e-2, rtol=1e-2) - + # Without affine res_no_affine = turbodiffusion_ops.layernorm(x, None, None, eps, elementwise_affine=False) if isinstance(res_no_affine, tuple): @@ -232,7 +234,7 @@ class TestTurboDiffusion: out_triton_no_affine = out_triton_no_affine.to(dtype) ln_no_affine = torch.nn.LayerNorm(dim, eps=eps, elementwise_affine=False, dtype=dtype).cuda() out_ref_no_affine = ln_no_affine(x) - + torch.testing.assert_close(out_triton_no_affine, out_ref_no_affine, atol=1e-2, rtol=1e-2) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -243,52 +245,51 @@ class TestTurboDiffusion: eps = 1e-5 w = torch.randn(dim, dtype=dtype, device="cuda", requires_grad=True) b = torch.randn(dim, dtype=dtype, device="cuda", requires_grad=True) - + # Triton implementation via Function y = turbodiffusion_ops.LayerNormFunction.apply(x, w, b, eps, True) loss = y.sum() loss.backward() - + x_grad = x.grad w_grad = w.grad b_grad = b.grad - + # Reference x_ref = x.detach().clone().requires_grad_() ln = torch.nn.LayerNorm(dim, eps=eps, elementwise_affine=True, dtype=dtype).cuda() ln.weight.data.copy_(w.detach()) ln.bias.data.copy_(b.detach()) - + y_ref = ln(x_ref) loss_ref = y_ref.sum() loss_ref.backward() - + torch.testing.assert_close(x_grad, x_ref.grad, atol=1e-2, rtol=1e-2) torch.testing.assert_close(w_grad, ln.weight.grad, atol=1e-2, rtol=1e-2) torch.testing.assert_close(b_grad, ln.bias.grad, atol=1e-2, rtol=1e-2) - @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_layer_norm_cuda(self, dtype): if fastvideo_kernel_ops is None or not hasattr(fastvideo_kernel_ops, "layer_norm_cuda"): pytest.skip("layer_norm_cuda not available") - + shape = (16, 128) x = torch.randn(shape, dtype=dtype, device="cuda") dim = shape[-1] eps = 1e-5 w = torch.randn(dim, dtype=dtype, device="cuda") b = torch.randn(dim, dtype=dtype, device="cuda") - + # C++ implementation # Signature: layer_norm_cuda(Input, eps, W, B, Output) -> Output out_cuda = torch.empty_like(x) fastvideo_kernel_ops.layer_norm_cuda(x, eps, w, b, out_cuda) - + # Reference ln = torch.nn.LayerNorm(dim, eps=eps, elementwise_affine=True, dtype=dtype).cuda() ln.weight.data.copy_(w) ln.bias.data.copy_(b) out_ref = ln(x) - + torch.testing.assert_close(out_cuda, out_ref, atol=1e-2, rtol=1e-2) diff --git a/fastvideo-kernel/tests/test_vmoba_correctness.py b/fastvideo-kernel/tests/test_vmoba_correctness.py index ff5db6091..2f31bafd3 100644 --- a/fastvideo-kernel/tests/test_vmoba_correctness.py +++ b/fastvideo-kernel/tests/test_vmoba_correctness.py @@ -5,6 +5,7 @@ import pytest import random from fastvideo_kernel import moba_attn_varlen + def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, device="cuda"): """ Generates random data for testing the variable-length attention function. @@ -21,7 +22,7 @@ def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, dev remaining_len = total_seqlen - sum(seqlens) if remaining_len > 0: seqlens.append(remaining_len) - else: # Adjust if sum exceeds total_seqlen + else: # Adjust if sum exceeds total_seqlen seqlens.append(avg_seqlen) current_sum = sum(seqlens) seqlens[-1] -= (current_sum - total_seqlen) @@ -52,9 +53,8 @@ def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, dev @pytest.mark.parametrize("select_mode", ["topk", "threshold"]) @pytest.mark.parametrize("threshold_type", ["query_head", "head_global", "overall"]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_moba_attn_varlen_forward( - batch_size, total_seqlen, num_heads, head_dim, moba_chunk_size, moba_topk, select_mode, threshold_type, dtype -): +def test_moba_attn_varlen_forward(batch_size, total_seqlen, num_heads, head_dim, moba_chunk_size, moba_topk, + select_mode, threshold_type, dtype): """ Tests the forward pass of moba_attn_varlen for basic correctness. It checks output shape, dtype, and for the presence of NaNs/Infs. @@ -62,9 +62,7 @@ def test_moba_attn_varlen_forward( if dtype == torch.float32: pytest.skip("float32 is not supported in flash attention") - q, k, v, cu_seqlens, max_seqlen = generate_test_data( - batch_size, total_seqlen, num_heads, head_dim, dtype - ) + q, k, v, cu_seqlens, max_seqlen = generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype) # Ensure chunk size is not larger than the smallest sequence length min_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).min().item() @@ -82,7 +80,7 @@ def test_moba_attn_varlen_forward( moba_topk=moba_topk, select_mode=select_mode, threshold_type=threshold_type, - simsum_threshold=0.5, # A reasonable default for threshold mode + simsum_threshold=0.5, # A reasonable default for threshold mode ) except Exception as e: pytest.fail(f"moba_attn_varlen forward pass failed with exception: {e}") diff --git a/fastvideo-kernel/tests/test_vsa.py b/fastvideo-kernel/tests/test_vsa.py index 0f7dd1a69..85bb85e43 100644 --- a/fastvideo-kernel/tests/test_vsa.py +++ b/fastvideo-kernel/tests/test_vsa.py @@ -12,13 +12,14 @@ from .utils import generate_block_sparse_mask_for_function, create_full_mask_fro BLOCK_M = 64 BLOCK_N = 64 + def pytorch_test(Q, K, V, block_sparse_mask, dO): q_ = Q.clone().float().requires_grad_() k_ = K.clone().float().requires_grad_() v_ = V.clone().float().requires_grad_() QK = torch.matmul(q_, k_.transpose(-2, -1)) - QK /= (q_.size(-1) ** 0.5) + QK /= (q_.size(-1)**0.5) QK = QK.masked_fill(~block_sparse_mask.unsqueeze(0), float('-inf')) QK = torch.nn.functional.softmax(QK, dim=-1) @@ -34,19 +35,18 @@ def pytorch_test(Q, K, V, block_sparse_mask, dO): ) -def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, q_non_pad_index, kv_non_pad_index, q_num_blocks, kv_num_blocks, dO): +def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, q_non_pad_index, kv_non_pad_index, + q_num_blocks, kv_num_blocks, dO): Q = Q.detach().requires_grad_() K = K.detach().requires_grad_() V = V.detach().requires_grad_() - + q_padded = vsa_pad(Q, q_non_pad_index, q_num_blocks, BLOCK_M) k_padded = vsa_pad(K, kv_non_pad_index, kv_num_blocks, BLOCK_M) v_padded = vsa_pad(V, kv_non_pad_index, kv_num_blocks, BLOCK_M) # Use autograd-enabled wrapper (internally dispatches to SM90 kernel or Triton) from fastvideo_kernel.block_sparse_attn import block_sparse_attn - output_padded, _aux = block_sparse_attn( - q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes - ) + output_padded, _aux = block_sparse_attn(q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes) output = output_padded[:, :, q_non_pad_index, :] output.backward(dO) @@ -65,12 +65,14 @@ def get_non_pad_index( return index_pad[index_mask] + def generate_tensor(shape, dtype, device): tensor = torch.randn(shape, dtype=dtype, device=device) return tensor + def generate_variable_block_sizes(num_blocks, min_size=16, max_size=64, device="cuda"): - return torch.randint(min_size, max_size + 1, (num_blocks,), device=device, dtype=torch.int32) + return torch.randint(min_size, max_size + 1, (num_blocks, ), device=device, dtype=torch.int32) def vsa_pad(x, non_pad_index, num_blocks, block_size): @@ -78,12 +80,29 @@ def vsa_pad(x, non_pad_index, num_blocks, block_size): padded_x[:, :, non_pad_index, :] = x return padded_x -def check_correctness(h, d, num_blocks, k, num_iterations=20, error_mode='all'): + +def check_correctness(h, d, num_blocks, k, num_iterations=20, error_mode='all'): results = { - 'gO': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'gQ': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'gK': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'gV': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, + 'gO': { + 'sum_diff': 0.0, + 'sum_abs': 0.0, + 'max_diff': 0.0 + }, + 'gQ': { + 'sum_diff': 0.0, + 'sum_abs': 0.0, + 'max_diff': 0.0 + }, + 'gK': { + 'sum_diff': 0.0, + 'sum_abs': 0.0, + 'max_diff': 0.0 + }, + 'gV': { + 'sum_diff': 0.0, + 'sum_abs': 0.0, + 'max_diff': 0.0 + }, } device = "cuda" if torch.cuda.is_available() else "cpu" @@ -94,19 +113,21 @@ def check_correctness(h, d, num_blocks, k, num_iterations=20, error_mode='all') block_mask = generate_block_sparse_mask_for_function(h, num_blocks, num_blocks, k, device) full_mask = create_full_mask_from_block_mask(block_mask, variable_block_sizes, variable_block_sizes, device) for _ in range(num_iterations): - Q = generate_tensor((1, h, S, d), torch.bfloat16, device) + Q = generate_tensor((1, h, S, d), torch.bfloat16, device) K = generate_tensor((1, h, S, d), torch.bfloat16, device) - V = generate_tensor((1, h, S, d), torch.bfloat16, device) + V = generate_tensor((1, h, S, d), torch.bfloat16, device) dO = generate_tensor((1, h, S, d), torch.bfloat16, device) - + # print(Q.shape, K.shape, V.shape, dO.shape) # dO_padded = torch.zeros_like(dO_padded) # dO_padded[:, :, non_pad_index, :] = dO pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, full_mask, dO) - bs_o, bs_qg, bs_kg, bs_vg = block_sparse_kernel_test(Q, K, V, block_mask.unsqueeze(0), variable_block_sizes, non_pad_index, non_pad_index, num_blocks, num_blocks, dO) - for name, (pt, bs) in zip(['gQ', 'gK', 'gV', 'gO'], [(pt_qg, bs_qg), (pt_kg, bs_kg), (pt_vg, bs_vg), (pt_o, bs_o)]): + bs_o, bs_qg, bs_kg, bs_vg = block_sparse_kernel_test(Q, K, V, block_mask.unsqueeze(0), variable_block_sizes, + non_pad_index, non_pad_index, num_blocks, num_blocks, dO) + for name, (pt, bs) in zip(['gQ', 'gK', 'gV', 'gO'], [(pt_qg, bs_qg), (pt_kg, bs_kg), (pt_vg, bs_vg), + (pt_o, bs_o)]): if bs is not None: diff = pt - bs abs_diff = torch.abs(diff) @@ -125,39 +146,59 @@ def check_correctness(h, d, num_blocks, k, num_iterations=20, error_mode='all') return results + def check_correctness_qkdiff(h, d, num_q_blocks, num_kv_blocks, k, num_iterations=20, error_mode='all'): results = { - 'gO': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'gQ': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'gK': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'gV': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, + 'gO': { + 'sum_diff': 0.0, + 'sum_abs': 0.0, + 'max_diff': 0.0 + }, + 'gQ': { + 'sum_diff': 0.0, + 'sum_abs': 0.0, + 'max_diff': 0.0 + }, + 'gK': { + 'sum_diff': 0.0, + 'sum_abs': 0.0, + 'max_diff': 0.0 + }, + 'gV': { + 'sum_diff': 0.0, + 'sum_abs': 0.0, + 'max_diff': 0.0 + }, } device = "cuda" if torch.cuda.is_available() else "cpu" q_variable_block_sizes = generate_variable_block_sizes(num_q_blocks, device=device) kv_variable_block_sizes = generate_variable_block_sizes(num_kv_blocks, device=device) - + S_q = int(q_variable_block_sizes.sum().item()) S_kv = int(kv_variable_block_sizes.sum().item()) - + q_non_pad_index = get_non_pad_index(q_variable_block_sizes, num_q_blocks, BLOCK_M) kv_non_pad_index = get_non_pad_index(kv_variable_block_sizes, num_kv_blocks, BLOCK_M) - + block_mask = generate_block_sparse_mask_for_function(h, num_q_blocks, num_kv_blocks, k, device) full_mask = create_full_mask_from_block_mask(block_mask, q_variable_block_sizes, kv_variable_block_sizes, device) - + for _ in range(num_iterations): - Q = generate_tensor((1, h, S_q, d), torch.bfloat16, device) + Q = generate_tensor((1, h, S_q, d), torch.bfloat16, device) K = generate_tensor((1, h, S_kv, d), torch.bfloat16, device) - V = generate_tensor((1, h, S_kv, d), torch.bfloat16, device) + V = generate_tensor((1, h, S_kv, d), torch.bfloat16, device) dO = generate_tensor((1, h, S_q, d), torch.bfloat16, device) - + # print(Q.shape, K.shape, V.shape, dO.shape) pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, full_mask, dO) - bs_o, bs_qg, bs_kg, bs_vg = block_sparse_kernel_test(Q, K, V, block_mask.unsqueeze(0), kv_variable_block_sizes, q_non_pad_index, kv_non_pad_index, num_q_blocks, num_kv_blocks, dO) - - for name, (pt, bs) in zip(['gQ', 'gK', 'gV', 'gO'], [(pt_qg, bs_qg), (pt_kg, bs_kg), (pt_vg, bs_vg), (pt_o, bs_o)]): + bs_o, bs_qg, bs_kg, bs_vg = block_sparse_kernel_test(Q, K, V, block_mask.unsqueeze(0), kv_variable_block_sizes, + q_non_pad_index, kv_non_pad_index, num_q_blocks, + num_kv_blocks, dO) + + for name, (pt, bs) in zip(['gQ', 'gK', 'gV', 'gO'], [(pt_qg, bs_qg), (pt_kg, bs_kg), (pt_vg, bs_vg), + (pt_o, bs_o)]): if bs is not None: diff = pt - bs abs_diff = torch.abs(diff) @@ -170,7 +211,7 @@ def check_correctness_qkdiff(h, d, num_q_blocks, num_kv_blocks, k, num_iteration total_elements_q = h * S_q * d * num_iterations total_elements_kv = h * S_kv * d * num_iterations - + for name, data in results.items(): total_elements = total_elements_q if name in ['gQ', 'gO'] else total_elements_kv avg_diff = data['sum_diff'] / total_elements @@ -179,13 +220,26 @@ def check_correctness_qkdiff(h, d, num_q_blocks, num_kv_blocks, k, num_iteration return results + def generate_error_graphs(h, d, error_mode='all'): test_configs = [ - {"num_blocks": 16, "k": 2, "description": "Small sequence"}, - {"num_blocks": 32, "k": 4, "description": "Medium sequence"}, - {"num_blocks": 53, "k": 6, "description": "Large sequence"}, + { + "num_blocks": 16, + "k": 2, + "description": "Small sequence" + }, + { + "num_blocks": 32, + "k": 4, + "description": "Medium sequence" + }, + { + "num_blocks": 53, + "k": 6, + "description": "Large sequence" + }, ] - + print(f"\nError Analysis for h={h}, d={d}, mode={error_mode}") print("=" * 150) print(f"{'Config':<20} {'Blocks':<8} {'K':<4} " @@ -208,15 +262,41 @@ def generate_error_graphs(h, d, error_mode='all'): print("-" * 150) + def generate_error_graphs_qkdiff(h, d, error_mode='all'): test_configs = [ - {"num_q_blocks": 16, "num_kv_blocks": 32, "k": 2, "description": "Small Q, Med KV"}, - {"num_q_blocks": 32, "num_kv_blocks": 16, "k": 4, "description": "Med Q, Small KV"}, - {"num_q_blocks": 53, "num_kv_blocks": 32, "k": 6, "description": "Large Q, Med KV"}, - {"num_q_blocks": 16, "num_kv_blocks": 48, "k": 2, "description": "Small Q, Large KV"}, - {"num_q_blocks": 48, "num_kv_blocks": 16, "k": 2, "description": "Large Q, Small KV"}, + { + "num_q_blocks": 16, + "num_kv_blocks": 32, + "k": 2, + "description": "Small Q, Med KV" + }, + { + "num_q_blocks": 32, + "num_kv_blocks": 16, + "k": 4, + "description": "Med Q, Small KV" + }, + { + "num_q_blocks": 53, + "num_kv_blocks": 32, + "k": 6, + "description": "Large Q, Med KV" + }, + { + "num_q_blocks": 16, + "num_kv_blocks": 48, + "k": 2, + "description": "Small Q, Large KV" + }, + { + "num_q_blocks": 48, + "num_kv_blocks": 16, + "k": 2, + "description": "Large Q, Small KV" + }, ] - + print(f"\nError Analysis (QK Diff) for h={h}, d={d}, mode={error_mode}") print("=" * 150) print(f"{'Config':<20} {'Q Blks':<8} {'KV Blks':<8} {'K':<4} " @@ -240,6 +320,7 @@ def generate_error_graphs_qkdiff(h, d, error_mode='all'): print("-" * 150) + def test_video_sparse_attention_backward(): if not torch.cuda.is_available(): return @@ -252,6 +333,6 @@ def test_video_sparse_attention_backward(): generate_error_graphs_qkdiff(h, d, error_mode=mode) print("\nAnalysis completed for all modes.") + if __name__ == "__main__": test_video_sparse_attention_backward() - diff --git a/fastvideo-kernel/tests/test_vsa256_forward.py b/fastvideo-kernel/tests/test_vsa256_forward.py index 666eb4303..d493c41bc 100644 --- a/fastvideo-kernel/tests/test_vsa256_forward.py +++ b/fastvideo-kernel/tests/test_vsa256_forward.py @@ -66,8 +66,8 @@ def _run_case( k = torch.randn(batch, heads, skv, head_dim, device=device, dtype=dtype) v = torch.randn(batch, heads, skv, head_dim, device=device, dtype=dtype) - q_var = torch.full((q_blocks_256,), q_block, dtype=torch.int32, device=device) - kv_var = torch.full((kv_blocks_256,), kv_block_logical, dtype=torch.int32, device=device) + q_var = torch.full((q_blocks_256, ), q_block, dtype=torch.int32, device=device) + kv_var = torch.full((kv_blocks_256, ), kv_block_logical, dtype=torch.int32, device=device) # Reproduce the compression branch's per-block average so we can pre-compute # the top-k mask that the kernel will see and feed it to the dense reference. @@ -78,21 +78,20 @@ def _run_case( k_c = (k_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(k.dtype) v_c = (v_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(v.dtype) - scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (head_dim ** 0.5) + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (head_dim**0.5) attn = torch.softmax(scores, dim=-1) out_c = torch.matmul(attn, v_c) - out_c = ( - out_c.view(batch, heads, q_blocks_256, 1, head_dim) - .repeat(1, 1, 1, q_block, 1) - .view(batch, heads, sq, head_dim) - ) + out_c = (out_c.view(batch, heads, q_blocks_256, 1, head_dim).repeat(1, 1, 1, q_block, + 1).view(batch, heads, sq, head_dim)) topk_idx = torch.topk(scores, topk_logical, dim=-1).indices mask_256 = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)[0] full_mask = create_full_mask_from_block_mask(mask_256, q_var, kv_var, device=device) out_ref = out_c + _dense_reference(q, k, v, full_mask) out = video_sparse_attn( - q, k, v, + q, + k, + v, kv_var, q_var, topk_logical, @@ -112,7 +111,11 @@ def test_vsa256_forward_qk_equal() -> None: if not torch.cuda.is_available(): pytest.skip("CUDA is required") avg_abs, max_rel = _run_case( - heads=8, head_dim=128, q_blocks_256=8, kv_blocks_256=8, topk_logical=2, + heads=8, + head_dim=128, + q_blocks_256=8, + kv_blocks_256=8, + topk_logical=2, ) print(f"[vsa256 qk_equal] avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}") assert avg_abs < 5e-2 @@ -124,7 +127,11 @@ def test_vsa256_forward_qk_diff() -> None: if not torch.cuda.is_available(): pytest.skip("CUDA is required") avg_abs, max_rel = _run_case( - heads=8, head_dim=128, q_blocks_256=8, kv_blocks_256=12, topk_logical=2, + heads=8, + head_dim=128, + q_blocks_256=8, + kv_blocks_256=12, + topk_logical=2, ) print(f"[vsa256 qk_diff] avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}") assert avg_abs < 5e-2 diff --git a/fastvideo-kernel/tests/test_vsa256_forward_cross.py b/fastvideo-kernel/tests/test_vsa256_forward_cross.py index 6527ab37b..7b75b9817 100644 --- a/fastvideo-kernel/tests/test_vsa256_forward_cross.py +++ b/fastvideo-kernel/tests/test_vsa256_forward_cross.py @@ -49,20 +49,11 @@ def _torch_vsa256_reference( scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / math.sqrt(dim) attn = torch.softmax(scores, dim=-1) out_c = torch.matmul(attn, v_c) - out_c = ( - out_c.view(bsz, heads, q_blocks, 1, dim) - .repeat(1, 1, 1, q_block, 1) - .view_as(q) - ) + out_c = (out_c.view(bsz, heads, q_blocks, 1, dim).repeat(1, 1, 1, q_block, 1).view_as(q)) topk_idx = torch.topk(scores, topk_logical, dim=-1).indices block_mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True) - token_mask = ( - block_mask - .repeat_interleave(q_block, dim=2) - .repeat_interleave(kv_block, dim=3) - .to(torch.bool) - ) + token_mask = (block_mask.repeat_interleave(q_block, dim=2).repeat_interleave(kv_block, dim=3).to(torch.bool)) qf, kf, vf = q.float(), k.float(), v.float() logits = torch.matmul(qf, kf.transpose(-2, -1)) / math.sqrt(dim) @@ -99,15 +90,19 @@ def test_vsa256_forward_cross_torch_cute_triton(monkeypatch) -> None: q = torch.randn(bsz, heads, sq, dim, device=device, dtype=dtype) k = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype) v = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype) - q_var = torch.full((q_blocks_256,), q_block, dtype=torch.int32, device=device) - kv_var = torch.full((kv_blocks_256,), kv_block, dtype=torch.int32, device=device) + q_var = torch.full((q_blocks_256, ), q_block, dtype=torch.int32, device=device) + kv_var = torch.full((kv_blocks_256, ), kv_block, dtype=torch.int32, device=device) out_torch = _torch_vsa256_reference(q, k, v, q_var, kv_var, topk_logical) # CuTe (default). out_cute = video_sparse_attn( - q, k, v, - kv_var, q_var, topk_logical, + q, + k, + v, + kv_var, + q_var, + topk_logical, block_size=(4, 8, 8), compress_attn_weight=None, ) @@ -115,8 +110,12 @@ def test_vsa256_forward_cross_torch_cute_triton(monkeypatch) -> None: # Triton via route-A 256->64 expansion. monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1") out_triton = video_sparse_attn( - q, k, v, - kv_var, q_var, topk_logical, + q, + k, + v, + kv_var, + q_var, + topk_logical, block_size=(4, 8, 8), compress_attn_weight=None, ) @@ -127,12 +126,10 @@ def test_vsa256_forward_cross_torch_cute_triton(monkeypatch) -> None: torch_vs_cute = _metrics(out_torch, out_cute) torch_vs_triton = _metrics(out_torch, out_triton) cute_vs_triton = _metrics(out_cute, out_triton) - print( - "[cross-forward] " - f"torch_vs_cute(avg_abs={torch_vs_cute[0]:.6e}, max_rel={torch_vs_cute[1]:.6e}), " - f"torch_vs_triton(avg_abs={torch_vs_triton[0]:.6e}, max_rel={torch_vs_triton[1]:.6e}), " - f"cute_vs_triton(avg_abs={cute_vs_triton[0]:.6e}, max_rel={cute_vs_triton[1]:.6e})" - ) + print("[cross-forward] " + f"torch_vs_cute(avg_abs={torch_vs_cute[0]:.6e}, max_rel={torch_vs_cute[1]:.6e}), " + f"torch_vs_triton(avg_abs={torch_vs_triton[0]:.6e}, max_rel={torch_vs_triton[1]:.6e}), " + f"cute_vs_triton(avg_abs={cute_vs_triton[0]:.6e}, max_rel={cute_vs_triton[1]:.6e})") assert torch_vs_cute[0] < 1e-3 and torch_vs_cute[1] < 0.2 assert torch_vs_triton[0] < 1e-3 and torch_vs_triton[1] < 0.2 diff --git a/fastvideo-kernel/tests/test_vsa256_forward_vbs.py b/fastvideo-kernel/tests/test_vsa256_forward_vbs.py index e993f82cc..f3f59448a 100644 --- a/fastvideo-kernel/tests/test_vsa256_forward_vbs.py +++ b/fastvideo-kernel/tests/test_vsa256_forward_vbs.py @@ -48,25 +48,15 @@ def _torch_vsa256_reference( scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / math.sqrt(dim) attn = torch.softmax(scores, dim=-1) out_c = torch.matmul(attn, v_c) - out_c = ( - out_c.view(bsz, heads, q_blocks, 1, dim) - .repeat(1, 1, 1, q_block, 1) - .view_as(q) - ) + out_c = (out_c.view(bsz, heads, q_blocks, 1, dim).repeat(1, 1, 1, q_block, 1).view_as(q)) with torch.no_grad(): topk_idx = torch.topk(scores.detach(), topk_logical, dim=-1).indices block_mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True) block_token_idx = torch.arange(kv_block, device=kv_var.device, dtype=torch.int32) - kv_token_valid_by_block = ( - block_token_idx.view(1, -1) < kv_var.to(torch.int32).view(-1, 1) - ).to(torch.bool) + kv_token_valid_by_block = (block_token_idx.view(1, -1) < kv_var.to(torch.int32).view(-1, 1)).to(torch.bool) kv_token_valid = kv_token_valid_by_block.reshape(1, 1, 1, kv_blocks * kv_block) - token_mask = ( - block_mask - .repeat_interleave(q_block, dim=2) - .repeat_interleave(kv_block, dim=3) - ) + token_mask = (block_mask.repeat_interleave(q_block, dim=2).repeat_interleave(kv_block, dim=3)) token_mask = token_mask & kv_token_valid qf, kf, vf = q.float(), k.float(), v.float() @@ -98,8 +88,8 @@ def test_vsa256_cute_variable_block_size_vs_torch_ref() -> None: k = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype) v = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype) - q_var = torch.full((q_blocks_256,), q_block, dtype=torch.int32, device=device) - kv_var = torch.randint(16, kv_block + 1, (kv_blocks_256,), dtype=torch.int32, device=device) + q_var = torch.full((q_blocks_256, ), q_block, dtype=torch.int32, device=device) + kv_var = torch.randint(16, kv_block + 1, (kv_blocks_256, ), dtype=torch.int32, device=device) token_idx = torch.arange(kv_block, device=device, dtype=torch.int32) kv_valid = token_idx.view(1, -1) < kv_var.view(-1, 1) kv_valid = kv_valid.view(1, 1, kv_blocks_256, kv_block, 1) @@ -108,7 +98,9 @@ def test_vsa256_cute_variable_block_size_vs_torch_ref() -> None: v = v * kv_valid.to(v.dtype) out = video_sparse_attn( - q, k, v, + q, + k, + v, kv_var, q_var, topk_logical, @@ -120,10 +112,8 @@ def test_vsa256_cute_variable_block_size_vs_torch_ref() -> None: diff = (out - out_ref).abs() avg_abs = diff.mean().item() max_rel = (diff.max() / (out_ref.abs().mean() + 1e-6)).item() - print( - f"[vsa256-cute-vbs] kv_var[min={int(kv_var.min().item())}, " - f"max={int(kv_var.max().item())}], " - f"avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}" - ) + print(f"[vsa256-cute-vbs] kv_var[min={int(kv_var.min().item())}, " + f"max={int(kv_var.max().item())}], " + f"avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}") assert avg_abs < 1e-3 and max_rel < 0.2 diff --git a/fastvideo-kernel/tests/test_vsa256_triton.py b/fastvideo-kernel/tests/test_vsa256_triton.py index 400953785..53688b171 100644 --- a/fastvideo-kernel/tests/test_vsa256_triton.py +++ b/fastvideo-kernel/tests/test_vsa256_triton.py @@ -39,25 +39,15 @@ def _torch_vsa256_reference( scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / math.sqrt(dim) attn = torch.softmax(scores, dim=-1) out_c = torch.matmul(attn, v_c) - out_c = ( - out_c.view(bsz, heads, q_blocks, 1, dim) - .repeat(1, 1, 1, q_block, 1) - .view_as(q) - ) + out_c = (out_c.view(bsz, heads, q_blocks, 1, dim).repeat(1, 1, 1, q_block, 1).view_as(q)) with torch.no_grad(): topk_idx = torch.topk(scores.detach(), topk_logical, dim=-1).indices block_mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True) block_token_idx = torch.arange(kv_block, device=kv_var.device, dtype=torch.int32) - kv_token_valid_by_block = ( - block_token_idx.view(1, -1) < kv_var.to(torch.int32).view(-1, 1) - ).to(torch.bool) + kv_token_valid_by_block = (block_token_idx.view(1, -1) < kv_var.to(torch.int32).view(-1, 1)).to(torch.bool) kv_token_valid = kv_token_valid_by_block.reshape(1, 1, 1, kv_blocks * kv_block) - token_mask = ( - block_mask - .repeat_interleave(q_block, dim=2) - .repeat_interleave(kv_block, dim=3) - ) + token_mask = (block_mask.repeat_interleave(q_block, dim=2).repeat_interleave(kv_block, dim=3)) token_mask = token_mask & kv_token_valid qf, kf, vf = q.float(), k.float(), v.float() @@ -98,8 +88,8 @@ def test_vsa256_triton_forward_backward_vs_torch_ref(monkeypatch) -> None: v_base = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype) grad_out = torch.randn_like(q_base) - q_var = torch.full((q_blocks_256,), q_block, dtype=torch.int32, device=device) - kv_var = torch.randint(16, kv_block + 1, (kv_blocks_256,), dtype=torch.int32, device=device) + q_var = torch.full((q_blocks_256, ), q_block, dtype=torch.int32, device=device) + kv_var = torch.randint(16, kv_block + 1, (kv_blocks_256, ), dtype=torch.int32, device=device) token_idx = torch.arange(kv_block, device=device, dtype=torch.int32) kv_valid = token_idx.view(1, -1) < kv_var.view(-1, 1) kv_valid = kv_valid.view(1, 1, kv_blocks_256, kv_block, 1) @@ -111,7 +101,9 @@ def test_vsa256_triton_forward_backward_vs_torch_ref(monkeypatch) -> None: k = k_base.detach().clone().requires_grad_(True) v = v_base.detach().clone().requires_grad_(True) out = video_sparse_attn( - q, k, v, + q, + k, + v, kv_var, q_var, topk_logical, @@ -135,14 +127,12 @@ def test_vsa256_triton_forward_backward_vs_torch_ref(monkeypatch) -> None: m_dq = _metrics(dq_ref, dq) m_dk = _metrics(dk_ref, dk) m_dv = _metrics(dv_ref, dv) - print( - "[vsa256-triton] " - f"kv_var[min={int(kv_var.min().item())}, max={int(kv_var.max().item())}], " - f"out(avg_abs={m_out[0]:.6e}, max_rel={m_out[1]:.6e}), " - f"dq(avg_abs={m_dq[0]:.6e}, max_rel={m_dq[1]:.6e}), " - f"dk(avg_abs={m_dk[0]:.6e}, max_rel={m_dk[1]:.6e}), " - f"dv(avg_abs={m_dv[0]:.6e}, max_rel={m_dv[1]:.6e})" - ) + print("[vsa256-triton] " + f"kv_var[min={int(kv_var.min().item())}, max={int(kv_var.max().item())}], " + f"out(avg_abs={m_out[0]:.6e}, max_rel={m_out[1]:.6e}), " + f"dq(avg_abs={m_dq[0]:.6e}, max_rel={m_dq[1]:.6e}), " + f"dk(avg_abs={m_dk[0]:.6e}, max_rel={m_dk[1]:.6e}), " + f"dv(avg_abs={m_dv[0]:.6e}, max_rel={m_dv[1]:.6e})") assert m_out[0] < 1e-3 and m_out[1] < 0.2 assert m_dq[0] < 2e-2 and m_dq[1] < 0.5 diff --git a/fastvideo-kernel/tests/test_vsa_forward.py b/fastvideo-kernel/tests/test_vsa_forward.py index 55ce8d14c..b7406d4d1 100644 --- a/fastvideo-kernel/tests/test_vsa_forward.py +++ b/fastvideo-kernel/tests/test_vsa_forward.py @@ -12,6 +12,7 @@ from .utils import ( from .test_vsa import BLOCK_M # Import from local test_vsa from . import test_vsa as ref + def pytorch_forward( Q: torch.Tensor, K: torch.Tensor, @@ -29,7 +30,7 @@ def pytorch_forward( v = V.clone().float() attn = torch.matmul(q, k.transpose(-2, -1)) # [1, h, S_q, S_kv] - attn = attn / (q.size(-1) ** 0.5) + attn = attn / (q.size(-1)**0.5) attn = attn.masked_fill(~block_sparse_mask.unsqueeze(0), float("-inf")) attn = torch.nn.functional.softmax(attn, dim=-1) out = torch.matmul(attn, v) # [1, h, S_q, d] @@ -61,9 +62,7 @@ def block_sparse_forward_test( # Use autograd-enabled wrapper (internally dispatches SM90 C++ vs Triton) from fastvideo_kernel.block_sparse_attn import block_sparse_attn try: - out_padded, _aux = block_sparse_attn( - q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes - ) + out_padded, _aux = block_sparse_attn(q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes) except RuntimeError as e: pytest.skip(str(e)) @@ -86,22 +85,16 @@ def run_forward_equal_qk( assert torch.cuda.is_available(), "VSA kernels require CUDA" device = "cuda" - variable_block_sizes = ref.generate_variable_block_sizes( - num_blocks, device=device - ) + variable_block_sizes = ref.generate_variable_block_sizes(num_blocks, device=device) S = int(variable_block_sizes.sum().item()) - non_pad_index = ref.get_non_pad_index( - variable_block_sizes, num_blocks, BLOCK_M - ) + non_pad_index = ref.get_non_pad_index(variable_block_sizes, num_blocks, BLOCK_M) - block_mask = generate_block_sparse_mask_for_function( - h, num_blocks, num_blocks, k, device - ) - full_mask = create_full_mask_from_block_mask( - block_mask, variable_block_sizes, variable_block_sizes, device - ) + block_mask = generate_block_sparse_mask_for_function(h, num_blocks, num_blocks, k, device) + full_mask = create_full_mask_from_block_mask(block_mask, variable_block_sizes, variable_block_sizes, device) print(f"[qkequal] h: {h}, d: {d}, num_blocks: {num_blocks}, k: {k}") - print(f"[qkequal] variable_block_sizes: {variable_block_sizes}, non_pad_index: {non_pad_index.shape}, block_mask: {block_mask.shape}, full_mask: {full_mask.shape}") + print( + f"[qkequal] variable_block_sizes: {variable_block_sizes}, non_pad_index: {non_pad_index.shape}, block_mask: {block_mask.shape}, full_mask: {full_mask.shape}" + ) sum_diff = 0.0 sum_abs = 0.0 max_rel_diff = 0.0 @@ -153,29 +146,17 @@ def run_forward_qk_diff( device = "cuda" - q_variable_block_sizes = ref.generate_variable_block_sizes( - num_q_blocks, device=device - ) - kv_variable_block_sizes = ref.generate_variable_block_sizes( - num_kv_blocks, device=device - ) + q_variable_block_sizes = ref.generate_variable_block_sizes(num_q_blocks, device=device) + kv_variable_block_sizes = ref.generate_variable_block_sizes(num_kv_blocks, device=device) S_q = int(q_variable_block_sizes.sum().item()) S_kv = int(kv_variable_block_sizes.sum().item()) - q_non_pad_index = ref.get_non_pad_index( - q_variable_block_sizes, num_q_blocks, BLOCK_M - ) - kv_non_pad_index = ref.get_non_pad_index( - kv_variable_block_sizes, num_kv_blocks, BLOCK_M - ) + q_non_pad_index = ref.get_non_pad_index(q_variable_block_sizes, num_q_blocks, BLOCK_M) + kv_non_pad_index = ref.get_non_pad_index(kv_variable_block_sizes, num_kv_blocks, BLOCK_M) - block_mask = generate_block_sparse_mask_for_function( - h, num_q_blocks, num_kv_blocks, k, device - ) - full_mask = create_full_mask_from_block_mask( - block_mask, q_variable_block_sizes, kv_variable_block_sizes, device - ) + block_mask = generate_block_sparse_mask_for_function(h, num_q_blocks, num_kv_blocks, k, device) + full_mask = create_full_mask_from_block_mask(block_mask, q_variable_block_sizes, kv_variable_block_sizes, device) sum_diff = 0.0 sum_abs = 0.0 @@ -222,13 +203,9 @@ def test_video_sparse_attention_forward(): print("\nForward Block Sparse Attention Check (QK Different)") print("=" * 80) - avg_err_diff, max_rel_diff = run_forward_qk_diff( - h, d, num_q_blocks=32, num_kv_blocks=48, k=2 - ) - print( - f"QK diff: avg |ΔO| = {avg_err_diff:.6e}, max rel ΔO = {max_rel_diff:.6e}" - ) + avg_err_diff, max_rel_diff = run_forward_qk_diff(h, d, num_q_blocks=32, num_kv_blocks=48, k=2) + print(f"QK diff: avg |ΔO| = {avg_err_diff:.6e}, max rel ΔO = {max_rel_diff:.6e}") + if __name__ == "__main__": test_video_sparse_attention_forward() - diff --git a/fastvideo-kernel/tests/test_vsa_utils.py b/fastvideo-kernel/tests/test_vsa_utils.py index 35e6b21cf..65a4cfc42 100644 --- a/fastvideo-kernel/tests/test_vsa_utils.py +++ b/fastvideo-kernel/tests/test_vsa_utils.py @@ -58,7 +58,7 @@ class TestGetTilePartitionIndices: device = torch.device("cpu") idx = get_tile_partition_indices(dit_seq_shape, tile_size, device) n = math.prod(dit_seq_shape) - assert idx.shape == (n,) + assert idx.shape == (n, ) assert idx.dtype == torch.long assert set(idx.tolist()) == set(range(n)) @@ -73,7 +73,7 @@ class TestGetTilePartitionIndices: device = torch.device("cpu") shape = (5, 7, 3) idx = get_tile_partition_indices(shape, (4, 4, 4), device) - assert idx.shape == (5 * 7 * 3,) + assert idx.shape == (5 * 7 * 3, ) assert set(idx.tolist()) == set(range(5 * 7 * 3)) @@ -188,9 +188,12 @@ class TestBuildVsaMetadata: """build_vsa_metadata returns all expected keys.""" meta = build_vsa_metadata((8, 16, 16), device="cpu") expected_keys = { - "tile_partition_indices", "reverse_tile_partition_indices", - "variable_block_sizes", "non_pad_index", - "num_tiles", "max_block_size", + "tile_partition_indices", + "reverse_tile_partition_indices", + "variable_block_sizes", + "non_pad_index", + "num_tiles", + "max_block_size", } assert set(meta.keys()) == expected_keys @@ -227,8 +230,8 @@ class TestBuildVsaMetadata: shape = (8, 16, 16) meta = build_vsa_metadata(shape, device="cpu") n = math.prod(shape) - assert meta["tile_partition_indices"].shape == (n,) - assert meta["reverse_tile_partition_indices"].shape == (n,) + assert meta["tile_partition_indices"].shape == (n, ) + assert meta["reverse_tile_partition_indices"].shape == (n, ) assert meta["variable_block_sizes"].sum().item() == n assert meta["non_pad_index"].shape[0] == n @@ -243,16 +246,14 @@ class TestConsistencyWithFramework: def _skip_if_no_framework(self): try: from fastvideo.attention.backends.video_sparse_attn import ( - get_tile_partition_indices as fw_get_tile, - ) + get_tile_partition_indices as fw_get_tile, ) except ImportError: pytest.skip("fastvideo framework not installed") @pytest.mark.parametrize("shape", [(8, 16, 16), (9, 10, 7)]) def test_tile_indices_match(self, shape): from fastvideo.attention.backends.video_sparse_attn import ( - get_tile_partition_indices as fw_get_tile, - ) + get_tile_partition_indices as fw_get_tile, ) device = torch.device("cpu") tile_size = (4, 4, 4) ours = get_tile_partition_indices(shape, tile_size, device) @@ -262,8 +263,7 @@ class TestConsistencyWithFramework: @pytest.mark.parametrize("shape", [(8, 16, 16), (9, 10, 7)]) def test_variable_block_sizes_match(self, shape): from fastvideo.attention.backends.video_sparse_attn import ( - construct_variable_block_sizes as fw_construct_vbs, - ) + construct_variable_block_sizes as fw_construct_vbs, ) device = torch.device("cpu") tile_size = (4, 4, 4) num_tiles = tuple(math.ceil(s / t) for s, t in zip(shape, tile_size)) diff --git a/fastvideo-kernel/tests/test_vsa_varlen.py b/fastvideo-kernel/tests/test_vsa_varlen.py index 31231dc29..45589148e 100644 --- a/fastvideo-kernel/tests/test_vsa_varlen.py +++ b/fastvideo-kernel/tests/test_vsa_varlen.py @@ -24,10 +24,15 @@ from fastvideo_kernel.block_sparse_attn_varlen import block_sparse_attn_varlen def _reference_per_sequence( - q_list, k_list, v_list, - block_masks, vbs_list, - non_pad_q_list, non_pad_kv_list, - q_nblocks_list, kv_nblocks_list, + q_list, + k_list, + v_list, + block_masks, + vbs_list, + non_pad_q_list, + non_pad_kv_list, + q_nblocks_list, + kv_nblocks_list, ): """Run per-sequence block_sparse_attn and concat outputs.""" outs = [] @@ -38,7 +43,12 @@ def _reference_per_sequence( q2k_idx, q2k_num = _map_to_index(block_masks[i].unsqueeze(0)) o_pad, _ = block_sparse_attn_from_indices( - q_pad, k_pad, v_pad, q2k_idx, q2k_num, vbs_list[i], + q_pad, + k_pad, + v_pad, + q2k_idx, + q2k_num, + vbs_list[i], ) o = o_pad[:, :, non_pad_q_list[i], :] outs.append(o.squeeze(0).transpose(0, 1)) @@ -111,29 +121,41 @@ def _run_varlen_test( cu_kv.append(cu_kv[-1] + skv) ref_out = _reference_per_sequence( - q_list, k_list, v_list, - block_masks, vbs_list, - non_pad_q_list, non_pad_kv_list, - q_nblocks_list, kv_nblocks_list, + q_list, + k_list, + v_list, + block_masks, + vbs_list, + non_pad_q_list, + non_pad_kv_list, + q_nblocks_list, + kv_nblocks_list, ) q_packed = torch.cat( - [qi.squeeze(0).transpose(0, 1) for qi in q_list], dim=0, + [qi.squeeze(0).transpose(0, 1) for qi in q_list], + dim=0, ) k_packed = torch.cat( - [ki.squeeze(0).transpose(0, 1) for ki in k_list], dim=0, + [ki.squeeze(0).transpose(0, 1) for ki in k_list], + dim=0, ) v_packed = torch.cat( - [vi.squeeze(0).transpose(0, 1) for vi in v_list], dim=0, + [vi.squeeze(0).transpose(0, 1) for vi in v_list], + dim=0, ) cu_seqlens_q = torch.tensor(cu_q, dtype=torch.int32, device=device) cu_seqlens_kv = torch.tensor(cu_kv, dtype=torch.int32, device=device) varlen_out = block_sparse_attn_varlen( - q_packed, k_packed, v_packed, - cu_seqlens_q, cu_seqlens_kv, - q2k_idx_list, q2k_num_list, + q_packed, + k_packed, + v_packed, + cu_seqlens_q, + cu_seqlens_kv, + q2k_idx_list, + q2k_num_list, vbs_list, q_variable_block_sizes_list=q_vbs_for_varlen, ) @@ -215,9 +237,14 @@ class TestVSAVarlen: cu_kv = torch.tensor([0, skv], dtype=torch.int32, device=device) varlen_out = block_sparse_attn_varlen( - q_flat, k_flat, v_flat, - cu_q, cu_kv, - [q2k_idx], [q2k_num], [vbs_kv], + q_flat, + k_flat, + v_flat, + cu_q, + cu_kv, + [q2k_idx], + [q2k_num], + [vbs_kv], ) max_abs = (ref_flat - varlen_out).abs().max().item() @@ -307,7 +334,12 @@ def _run_varlen_backward_test( q2k_idx, q2k_num = _map_to_index(block_masks[i].unsqueeze(0)) o_pad, _ = block_sparse_attn_from_indices( - q_pad, k_pad, v_pad, q2k_idx, q2k_num, vbs_list[i], + q_pad, + k_pad, + v_pad, + q2k_idx, + q2k_num, + vbs_list[i], ) o = o_pad[:, :, non_pad_q_list[i], :] o_flat = o.squeeze(0).transpose(0, 1) @@ -326,22 +358,29 @@ def _run_varlen_backward_test( # --- Varlen backward --- q_packed = torch.cat( - [qi.squeeze(0).transpose(0, 1) for qi in q_list], dim=0, + [qi.squeeze(0).transpose(0, 1) for qi in q_list], + dim=0, ).detach().requires_grad_(True) k_packed = torch.cat( - [ki.squeeze(0).transpose(0, 1) for ki in k_list], dim=0, + [ki.squeeze(0).transpose(0, 1) for ki in k_list], + dim=0, ).detach().requires_grad_(True) v_packed = torch.cat( - [vi.squeeze(0).transpose(0, 1) for vi in v_list], dim=0, + [vi.squeeze(0).transpose(0, 1) for vi in v_list], + dim=0, ).detach().requires_grad_(True) cu_seqlens_q = torch.tensor(cu_q, dtype=torch.int32, device=device) cu_seqlens_kv = torch.tensor(cu_kv, dtype=torch.int32, device=device) varlen_out = block_sparse_attn_varlen( - q_packed, k_packed, v_packed, - cu_seqlens_q, cu_seqlens_kv, - q2k_idx_list, q2k_num_list, + q_packed, + k_packed, + v_packed, + cu_seqlens_q, + cu_seqlens_kv, + q2k_idx_list, + q2k_num_list, vbs_list, q_variable_block_sizes_list=q_vbs_list, ) @@ -363,9 +402,7 @@ def _run_varlen_backward_test( mean_abs = ref.abs().mean().item() max_rel = max_abs / (mean_abs + 1e-8) print(f" {name}: max_abs={max_abs:.4e}, max_rel={max_rel:.4e}") - assert max_rel < grad_rtol, ( - f"{name}: max relative error {max_rel:.4e} exceeds threshold {grad_rtol}" - ) + assert max_rel < grad_rtol, (f"{name}: max relative error {max_rel:.4e} exceeds threshold {grad_rtol}") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -413,9 +450,14 @@ class TestVSAVarlenBackward: cu_kv = torch.tensor([0, skv], dtype=torch.int32, device=device) out = block_sparse_attn_varlen( - q, k, v, - cu_q, cu_kv, - [q2k_idx], [q2k_num], [vbs_kv], + q, + k, + v, + cu_q, + cu_kv, + [q2k_idx], + [q2k_num], + [vbs_kv], q_variable_block_sizes_list=[vbs_q], ) diff --git a/fastvideo-kernel/tests/utils.py b/fastvideo-kernel/tests/utils.py index 3d9f15f45..f8ef80b02 100644 --- a/fastvideo-kernel/tests/utils.py +++ b/fastvideo-kernel/tests/utils.py @@ -1,5 +1,6 @@ import torch + def generate_block_sparse_mask_for_function(h, num_q_blocks, num_kv_blocks, k, device="cuda"): """ Generate block sparse mask of shape [h, num_q_blocks, num_kv_blocks]. @@ -23,8 +24,7 @@ def generate_block_sparse_mask_for_function(h, num_q_blocks, num_kv_blocks, k, d return block_sparse_mask -def create_full_mask_from_block_mask(block_sparse_mask, q_variable_block_sizes, - kv_variable_block_sizes, device="cuda"): +def create_full_mask_from_block_mask(block_sparse_mask, q_variable_block_sizes, kv_variable_block_sizes, device="cuda"): """ Convert block-level sparse mask to full attention mask. @@ -40,7 +40,7 @@ def create_full_mask_from_block_mask(block_sparse_mask, q_variable_block_sizes, h, num_q_blocks, num_kv_blocks = block_sparse_mask.shape total_q_seq_len = q_variable_block_sizes.sum().item() total_kv_seq_len = kv_variable_block_sizes.sum().item() - + q_cumsum = torch.cat([torch.tensor([0], device=device), q_variable_block_sizes.cumsum(dim=0)[:-1]]) kv_cumsum = torch.cat([torch.tensor([0], device=device), kv_variable_block_sizes.cumsum(dim=0)[:-1]]) @@ -57,4 +57,4 @@ def create_full_mask_from_block_mask(block_sparse_mask, q_variable_block_sizes, kv_end = kv_start + kv_variable_block_sizes[kv_block] full_mask[head, q_start:q_end, kv_start:kv_end] = True - return full_mask \ No newline at end of file + return full_mask diff --git a/fastvideo/dataset/__init__.py b/fastvideo/dataset/__init__.py index b82c65327..99ac81efc 100644 --- a/fastvideo/dataset/__init__.py +++ b/fastvideo/dataset/__init__.py @@ -2,13 +2,10 @@ from torchvision import transforms from torchvision.transforms import Lambda -from fastvideo.dataset.parquet_dataset_map_style import ( - build_parquet_map_style_dataloader) -from fastvideo.dataset.ltx2_precomputed_dataset import ( - build_ltx2_precomputed_dataloader, LTX2PrecomputedDataset) +from fastvideo.dataset.parquet_dataset_map_style import (build_parquet_map_style_dataloader) +from fastvideo.dataset.ltx2_precomputed_dataset import (build_ltx2_precomputed_dataloader, LTX2PrecomputedDataset) from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset, TextDataset -from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255, - TemporalRandomCrop) +from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255, TemporalRandomCrop) from fastvideo.dataset.validation_dataset import ValidationDataset @@ -36,12 +33,10 @@ def getdataset(args) -> VideoCaptionMergedDataset: temporal_sample=temporal_sample, transform_topcrop=transform_topcrop, seed=args.seed) - + def gettextdataset(args) -> TextDataset: - return TextDataset(data_merge_path=args.data_merge_path, - args=args, - seed=args.seed) + return TextDataset(data_merge_path=args.data_merge_path, args=args, seed=args.seed) __all__ = [ diff --git a/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py b/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py index 893d81459..1062f49b6 100644 --- a/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py +++ b/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py @@ -7,39 +7,26 @@ import time import torch.distributed as dist import torch.distributed.checkpoint as dist_cp -from fastvideo.dataset.parquet_dataset_iterable_style import ( - build_parquet_iterable_style_dataloader) +from fastvideo.dataset.parquet_dataset_iterable_style import (build_parquet_iterable_style_dataloader) from fastvideo.distributed import get_world_rank -from fastvideo.distributed.parallel_state import ( - cleanup_dist_env_and_memory, get_local_torch_device, - maybe_init_distributed_environment_and_model_parallel) +from fastvideo.distributed.parallel_state import (cleanup_dist_env_and_memory, get_local_torch_device, + maybe_init_distributed_environment_and_model_parallel) from fastvideo.logger import init_logger logger = init_logger(__name__) def main() -> None: - parser = argparse.ArgumentParser( - description="Benchmark parquet iterable style dataset loading speed") + parser = argparse.ArgumentParser(description="Benchmark parquet iterable style dataset loading speed") parser.add_argument( "--path", type=str, help="Path to parquet dataset", ) - parser.add_argument("--batch_size", - type=int, - default=4, - help="Batch size for DataLoader") - parser.add_argument("--num_data_workers", - type=int, - help="Number of DataLoader workers") - parser.add_argument("--num_epoch", - type=int, - default=2, - help="Number of epoches to benchmark") - parser.add_argument("--verify_resume", - action="store_true", - help="Verify resume") + parser.add_argument("--batch_size", type=int, default=4, help="Batch size for DataLoader") + parser.add_argument("--num_data_workers", type=int, help="Number of DataLoader workers") + parser.add_argument("--num_epoch", type=int, default=2, help="Number of epoches to benchmark") + parser.add_argument("--verify_resume", action="store_true", help="Verify resume") parser.add_argument( "--num_batches_per_epoch", type=int, @@ -58,21 +45,17 @@ def main() -> None: ''' args = parser.parse_args() world_size = int(os.environ.get("WORLD_SIZE", 1)) - maybe_init_distributed_environment_and_model_parallel( - tp_size=(world_size + 1) // 2, sp_size=(world_size + 1) // 2) - logger.info("Initialized distributed environment with world_size=%d", - world_size) + maybe_init_distributed_environment_and_model_parallel(tp_size=(world_size + 1) // 2, sp_size=(world_size + 1) // 2) + logger.info("Initialized distributed environment with world_size=%d", world_size) # Create DataLoader with proper settings - dataset, dataloader = build_parquet_iterable_style_dataloader( - args.path, args.batch_size, args.num_data_workers) + dataset, dataloader = build_parquet_iterable_style_dataloader(args.path, args.batch_size, args.num_data_workers) logger.info("Initialized dataloader") if args.verify_resume: # First pass - record latent sums first_pass_sums = [] - for i, (latents, embeddings, masks, - caption_text) in enumerate(dataloader): + for i, (latents, embeddings, masks, caption_text) in enumerate(dataloader): latent_sum = latents.sum().item() first_pass_sums.append(latent_sum) logger.info("Batch %d latent sum: %f", i, latent_sum) @@ -81,69 +64,54 @@ def main() -> None: # Save dataloader state using distributed checkpoint checkpoint_dir = pathlib.Path(args.checkpoint_path) - logger.info("Rank %d: Saving dataloader state to %s", get_world_rank(), - checkpoint_dir) + logger.info("Rank %d: Saving dataloader state to %s", get_world_rank(), checkpoint_dir) states = {"dataloader": dataloader} begin_time = time.monotonic() dist_cp.save(states, checkpoint_id=checkpoint_dir.as_posix()) end_time = time.monotonic() - logger.info("Rank %d: Saved checkpoint in %.2f seconds", - get_world_rank(), end_time - begin_time) + logger.info("Rank %d: Saved checkpoint in %.2f seconds", get_world_rank(), end_time - begin_time) # Make sure all processes wait for checkpoint to be saved if world_size > 1: dist.barrier() # Recreate dataloader and load state - dataset, dataloader = build_parquet_iterable_style_dataloader( - args.path, args.batch_size, args.num_data_workers) + dataset, dataloader = build_parquet_iterable_style_dataloader(args.path, args.batch_size, args.num_data_workers) load_states = {"dataloader": dataloader} dist_cp.load(load_states, checkpoint_id=checkpoint_dir.as_posix()) - logger.info("Rank %d: Loaded dataloader state from %s", - get_world_rank(), checkpoint_dir) + logger.info("Rank %d: Loaded dataloader state from %s", get_world_rank(), checkpoint_dir) # Second pass - verify latent sums match for i, (latents, embeddings, masks) in enumerate(dataloader): latent_sum = latents.sum().item() first_pass_sums.append(latent_sum) - logger.info("Batch %d latent sum: %f", - i + args.num_batches_per_epoch, latent_sum) + logger.info("Batch %d latent sum: %f", i + args.num_batches_per_epoch, latent_sum) if i >= args.num_batches_per_epoch - 1: break - dataset, dataloader = build_parquet_iterable_style_dataloader( - args.path, args.batch_size, args.num_data_workers) + dataset, dataloader = build_parquet_iterable_style_dataloader(args.path, args.batch_size, args.num_data_workers) # Second pass - verify latent sums match second_pass_sums = [] - for i, (latents, embeddings, masks, - caption_text) in enumerate(dataloader): + for i, (latents, embeddings, masks, caption_text) in enumerate(dataloader): latent_sum = latents.sum().item() second_pass_sums.append(latent_sum) - logger.info("Batch %d latent sum: %f (should match first pass: %f)", - i, latent_sum, first_pass_sums[i]) + logger.info("Batch %d latent sum: %f (should match first pass: %f)", i, latent_sum, first_pass_sums[i]) if i >= args.num_batches_per_epoch * 2 - 1: break # Verify all sums match - if all( - abs(a - b) < 1e-6 for a, b in zip( - first_pass_sums, second_pass_sums, strict=True)): - logger.info( - "All latent sums match between passes - resume verification successful!" - ) + if all(abs(a - b) < 1e-6 for a, b in zip(first_pass_sums, second_pass_sums, strict=True)): + logger.info("All latent sums match between passes - resume verification successful!") else: - raise ValueError( - "Latent sums do not match between passes - resume verification failed!" - ) + raise ValueError("Latent sums do not match between passes - resume verification failed!") start_time = time.time() total_samples = 0 total_batches = 0 for _ in range(args.num_epoch): - for i, (latents, embeddings, masks, - caption_text) in enumerate(dataloader): + for i, (latents, embeddings, masks, caption_text) in enumerate(dataloader): if i >= args.num_batches_per_epoch: break @@ -160,8 +128,7 @@ def main() -> None: if get_world_rank() == 0 and (i + 1) % 10 == 0: elapsed = time.time() - start_time samples_per_sec = total_samples / elapsed - logger.info("Batch %d/%d, Speed: %.2f samples/sec", i + 1, - args.num_batches_per_epoch, samples_per_sec) + logger.info("Batch %d/%d, Speed: %.2f samples/sec", i + 1, args.num_batches_per_epoch, samples_per_sec) # Final statistics if world_size > 1: diff --git a/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_map_style.py b/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_map_style.py index 9413099c7..9ada11e73 100644 --- a/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_map_style.py +++ b/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_map_style.py @@ -9,12 +9,10 @@ import torch.distributed as dist import torch.distributed.checkpoint as dist_cp from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v -from fastvideo.dataset.parquet_dataset_map_style import ( - build_parquet_map_style_dataloader) +from fastvideo.dataset.parquet_dataset_map_style import (build_parquet_map_style_dataloader) from fastvideo.distributed import get_world_rank -from fastvideo.distributed.parallel_state import ( - cleanup_dist_env_and_memory, get_local_torch_device, - maybe_init_distributed_environment_and_model_parallel) +from fastvideo.distributed.parallel_state import (cleanup_dist_env_and_memory, get_local_torch_device, + maybe_init_distributed_environment_and_model_parallel) from fastvideo.logger import init_logger logger = init_logger(__name__) @@ -22,27 +20,16 @@ logger = init_logger(__name__) def main() -> None: torch.multiprocessing.set_start_method("spawn", force=True) - parser = argparse.ArgumentParser( - description="Benchmark parquet map style dataset loading speed") + parser = argparse.ArgumentParser(description="Benchmark parquet map style dataset loading speed") parser.add_argument( "--path", type=str, help="Path to parquet dataset", ) - parser.add_argument("--batch_size", - type=int, - default=4, - help="Batch size for DataLoader") - parser.add_argument("--num_data_workers", - type=int, - help="Number of DataLoader workers") - parser.add_argument("--num_epoch", - type=int, - default=2, - help="Number of epoches to benchmark") - parser.add_argument("--verify_resume", - action="store_true", - help="Verify resume") + parser.add_argument("--batch_size", type=int, default=4, help="Batch size for DataLoader") + parser.add_argument("--num_data_workers", type=int, help="Number of DataLoader workers") + parser.add_argument("--num_epoch", type=int, default=2, help="Number of epoches to benchmark") + parser.add_argument("--verify_resume", action="store_true", help="Verify resume") parser.add_argument( "--num_batches_per_epoch", type=int, @@ -61,17 +48,14 @@ def main() -> None: ''' args = parser.parse_args() world_size = int(os.environ.get("WORLD_SIZE", 1)) - maybe_init_distributed_environment_and_model_parallel( - tp_size=(world_size + 1) // 2, sp_size=(world_size + 1) // 2) - logger.info("Initialized distributed environment with world_size=%d", - world_size) + maybe_init_distributed_environment_and_model_parallel(tp_size=(world_size + 1) // 2, sp_size=(world_size + 1) // 2) + logger.info("Initialized distributed environment with world_size=%d", world_size) # Create DataLoader with proper settings - dataset, dataloader = build_parquet_map_style_dataloader( - args.path, - args.batch_size, - parquet_schema=pyarrow_schema_t2v, - num_data_workers=args.num_data_workers) + dataset, dataloader = build_parquet_map_style_dataloader(args.path, + args.batch_size, + parquet_schema=pyarrow_schema_t2v, + num_data_workers=args.num_data_workers) logger.info("Initialized dataloader with %d batches", len(dataloader)) if args.verify_resume: @@ -88,47 +72,41 @@ def main() -> None: # Save dataloader state using distributed checkpoint checkpoint_dir = pathlib.Path(args.checkpoint_path) - logger.info("Rank %d: Saving dataloader state to %s", get_world_rank(), - checkpoint_dir) + logger.info("Rank %d: Saving dataloader state to %s", get_world_rank(), checkpoint_dir) states = {"dataloader": dataloader} begin_time = time.monotonic() dist_cp.save(states, checkpoint_id=checkpoint_dir.as_posix()) end_time = time.monotonic() - logger.info("Rank %d: Saved checkpoint in %.2f seconds", - get_world_rank(), end_time - begin_time) + logger.info("Rank %d: Saved checkpoint in %.2f seconds", get_world_rank(), end_time - begin_time) # Make sure all processes wait for checkpoint to be saved if world_size > 1: dist.barrier() # Recreate dataloader and load state - dataset, dataloader = build_parquet_map_style_dataloader( - args.path, - args.batch_size, - parquet_schema=pyarrow_schema_t2v, - num_data_workers=args.num_data_workers) + dataset, dataloader = build_parquet_map_style_dataloader(args.path, + args.batch_size, + parquet_schema=pyarrow_schema_t2v, + num_data_workers=args.num_data_workers) load_states = {"dataloader": dataloader} dist_cp.load(load_states, checkpoint_id=checkpoint_dir.as_posix()) - logger.info("Rank %d: Loaded dataloader state from %s", - get_world_rank(), checkpoint_dir) + logger.info("Rank %d: Loaded dataloader state from %s", get_world_rank(), checkpoint_dir) for i, batch in enumerate(dataloader): latents = batch['vae_latent'] embeddings = batch['text_embedding'] latent_sum = latents.sum().item() first_pass_sums.append(latent_sum) - logger.info("Batch %d latent sum: %f", - i + args.num_batches_per_epoch, latent_sum) + logger.info("Batch %d latent sum: %f", i + args.num_batches_per_epoch, latent_sum) if i >= args.num_batches_per_epoch - 1: break - dataset, dataloader = build_parquet_map_style_dataloader( - args.path, - args.batch_size, - parquet_schema=pyarrow_schema_t2v, - num_data_workers=args.num_data_workers) + dataset, dataloader = build_parquet_map_style_dataloader(args.path, + args.batch_size, + parquet_schema=pyarrow_schema_t2v, + num_data_workers=args.num_data_workers) # Second pass - verify latent sums match second_pass_sums = [] @@ -137,22 +115,15 @@ def main() -> None: embeddings = batch['text_embedding'] latent_sum = latents.sum().item() second_pass_sums.append(latent_sum) - logger.info("Batch %d latent sum: %f (should match first pass: %f)", - i, latent_sum, first_pass_sums[i]) + logger.info("Batch %d latent sum: %f (should match first pass: %f)", i, latent_sum, first_pass_sums[i]) if i >= args.num_batches_per_epoch * 2 - 1: break # Verify all sums match - if all( - abs(a - b) < 1e-6 for a, b in zip( - first_pass_sums, second_pass_sums, strict=True)): - logger.info( - "All latent sums match between passes - resume verification successful!" - ) + if all(abs(a - b) < 1e-6 for a, b in zip(first_pass_sums, second_pass_sums, strict=True)): + logger.info("All latent sums match between passes - resume verification successful!") else: - raise ValueError( - "Latent sums do not match between passes - resume verification failed!" - ) + raise ValueError("Latent sums do not match between passes - resume verification failed!") start_time = time.time() total_samples = 0 @@ -177,8 +148,7 @@ def main() -> None: if get_world_rank() == 0 and (i + 1) % 10 == 0: elapsed = time.time() - start_time samples_per_sec = total_samples / elapsed - logger.info("Batch %d/%d, Speed: %.2f samples/sec", i + 1, - args.num_batches_per_epoch, samples_per_sec) + logger.info("Batch %d/%d, Speed: %.2f samples/sec", i + 1, args.num_batches_per_epoch, samples_per_sec) # Final statistics if world_size > 1: diff --git a/fastvideo/dataset/dataloader/parquet_io.py b/fastvideo/dataset/dataloader/parquet_io.py index 134382b85..3c7ce81f8 100644 --- a/fastvideo/dataset/dataloader/parquet_io.py +++ b/fastvideo/dataset/dataloader/parquet_io.py @@ -161,12 +161,9 @@ class ParquetDatasetWriter: for file in files: if file.endswith('.parquet'): num_parquets += 1 - remainder_path = os.path.join(worker_dir, - f"data_chunk_{num_parquets}.parquet") + remainder_path = os.path.join(worker_dir, f"data_chunk_{num_parquets}.parquet") temp_path = remainder_path + '.tmp' - pq.write_table(remainder_table, - temp_path, - compression=self.compression) + pq.write_table(remainder_table, temp_path, compression=self.compression) if os.path.exists(remainder_path): os.remove(remainder_path) os.rename(temp_path, remainder_path) @@ -184,17 +181,15 @@ class ParquetDatasetWriter: start_chunk = worker_id * chunks_per_worker end_chunk = min((worker_id + 1) * chunks_per_worker, total_chunks) if start_chunk < end_chunk: - work_ranges.append( - ( - start_chunk, - end_chunk, - table_to_write, - worker_id, - self.out_dir, - self.samples_per_file, - self.compression, - ) - ) + work_ranges.append(( + start_chunk, + end_chunk, + table_to_write, + worker_id, + self.out_dir, + self.samples_per_file, + self.compression, + )) written_total = 0 if len(work_ranges) == 1: @@ -259,6 +254,3 @@ def _process_chunk_range(args: Any) -> int: raise return total_written - - - diff --git a/fastvideo/dataset/dataloader/record_schema.py b/fastvideo/dataset/dataloader/record_schema.py index 03c9f58df..d60c52ae6 100644 --- a/fastvideo/dataset/dataloader/record_schema.py +++ b/fastvideo/dataset/dataloader/record_schema.py @@ -31,37 +31,21 @@ def basic_t2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]: duration_val = float(batch.num_frames[idx]) / float(batch.fps[idx]) record = { - "id": - video_name, - "vae_latent_bytes": - batch.latents[idx].tobytes(), - "vae_latent_shape": - list(batch.latents[idx].shape), - "vae_latent_dtype": - str(batch.latents[idx].dtype), - "text_embedding_bytes": - batch.prompt_embeds[idx].tobytes(), - "text_embedding_shape": - list(batch.prompt_embeds[idx].shape), - "text_embedding_dtype": - str(batch.prompt_embeds[idx].dtype), - "file_name": - video_name, - "caption": - batch.prompt[idx], - "media_type": - "video", - "width": - int(width), - "height": - int(height), - "num_frames": - batch.latents[idx].shape[1] - if len(batch.latents[idx].shape) > 1 else 0, - "duration_sec": - duration_val, - "fps": - fps_val, + "id": video_name, + "vae_latent_bytes": batch.latents[idx].tobytes(), + "vae_latent_shape": list(batch.latents[idx].shape), + "vae_latent_dtype": str(batch.latents[idx].dtype), + "text_embedding_bytes": batch.prompt_embeds[idx].tobytes(), + "text_embedding_shape": list(batch.prompt_embeds[idx].shape), + "text_embedding_dtype": str(batch.prompt_embeds[idx].dtype), + "file_name": video_name, + "caption": batch.prompt[idx], + "media_type": "video", + "width": int(width), + "height": int(height), + "num_frames": batch.latents[idx].shape[1] if len(batch.latents[idx].shape) > 1 else 0, + "duration_sec": duration_val, + "fps": fps_val, } records.append(record) @@ -73,8 +57,7 @@ def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]: records = basic_t2v_record_creator(batch) if len(batch.image_embeds) > 0: - assert len(batch.image_embeds - ) == 1, "image embedding should be a single tensor" + assert len(batch.image_embeds) == 1, "image embedding should be a single tensor" image_embeds = batch.image_embeds[0] else: image_embeds = None @@ -97,12 +80,9 @@ def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]: if image_latent is not None: record.update({ - "first_frame_latent_bytes": - image_latent[idx].tobytes(), - "first_frame_latent_shape": - list(image_latent[idx].shape), - "first_frame_latent_dtype": - str(image_latent[idx].dtype), + "first_frame_latent_bytes": image_latent[idx].tobytes(), + "first_frame_latent_shape": list(image_latent[idx].shape), + "first_frame_latent_dtype": str(image_latent[idx].dtype), }) else: record.update({ @@ -127,10 +107,8 @@ def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]: return records -def ode_text_only_record_creator( - video_name: str, text_embedding: np.ndarray, caption: str, - trajectory_latents: np.ndarray, - trajectory_timesteps: np.ndarray) -> dict[str, Any]: +def ode_text_only_record_creator(video_name: str, text_embedding: np.ndarray, caption: str, + trajectory_latents: np.ndarray, trajectory_timesteps: np.ndarray) -> dict[str, Any]: """Create a text-only ODE trajectory record matching pyarrow_schema_ode_trajectory_text_only. Args: @@ -171,8 +149,7 @@ def ode_text_only_record_creator( return record -def text_only_record_creator(text_name: str, text_embedding: np.ndarray, - caption: str) -> dict[str, Any]: +def text_only_record_creator(text_name: str, text_embedding: np.ndarray, caption: str) -> dict[str, Any]: """Create a text-only record matching pyarrow_schema_text_only. Args: @@ -193,16 +170,15 @@ def text_only_record_creator(text_name: str, text_embedding: np.ndarray, return record -def matrixgame2_ode_record_creator( - video_name: str, - clip_feature: np.ndarray, - first_frame_latent: np.ndarray, - trajectory_latents: np.ndarray, - trajectory_timesteps: np.ndarray, - pil_image: np.ndarray | None = None, - keyboard_cond: np.ndarray | None = None, - mouse_cond: np.ndarray | None = None, - caption: str = "") -> dict[str, Any]: +def matrixgame2_ode_record_creator(video_name: str, + clip_feature: np.ndarray, + first_frame_latent: np.ndarray, + trajectory_latents: np.ndarray, + trajectory_timesteps: np.ndarray, + pil_image: np.ndarray | None = None, + keyboard_cond: np.ndarray | None = None, + mouse_cond: np.ndarray | None = None, + caption: str = "") -> dict[str, Any]: """Create a ODE trajectory record matching pyarrow_schema_matrixgame2_ode_trajectory. Args: @@ -297,4 +273,4 @@ def matrixgame2_ode_record_creator( "trajectory_timesteps_dtype": str(trajectory_timesteps.dtype), }) - return record \ No newline at end of file + return record diff --git a/fastvideo/dataset/dataloader/schema.py b/fastvideo/dataset/dataloader/schema.py index cfd5035a4..6fe807dbc 100644 --- a/fastvideo/dataset/dataloader/schema.py +++ b/fastvideo/dataset/dataloader/schema.py @@ -50,7 +50,6 @@ pyarrow_schema_i2v = pa.schema([ pa.field("fps", pa.float64()), ]) - pyarrow_schema_t2v = pa.schema([ pa.field("id", pa.string()), # --- Image/Video VAE latents --- @@ -80,7 +79,6 @@ pyarrow_schema_t2v = pa.schema([ pa.field("fps", pa.float64()), ]) - # One text-to-video-and-audio (T2VA) row owns synchronized video, audio, and text tensors so # collate_rows_from_parquet_schema cannot pair targets from different samples. # Each tensor uses the bytes/shape/dtype triplet that the collator discovers. @@ -110,7 +108,6 @@ pyarrow_schema_t2va = pa.schema([ pa.field("audio_sample_rate", pa.int64()), ]) - pyarrow_schema_ode_trajectory_text_only = pa.schema([ pa.field("id", pa.string()), # --- Text encoder output tensor --- @@ -133,7 +130,6 @@ pyarrow_schema_ode_trajectory_text_only = pa.schema([ pa.field("media_type", pa.string()), # Always 'text' for text-only ]) - pyarrow_schema_text_only = pa.schema([ pa.field("id", pa.string()), # --- Text encoder output tensor --- @@ -147,7 +143,6 @@ pyarrow_schema_text_only = pa.schema([ pa.field("caption", pa.string()), ]) - pyarrow_schema_matrixgame2 = pa.schema([ pa.field("id", pa.string()), # --- Image/Video VAE latents --- @@ -219,4 +214,4 @@ pyarrow_schema_matrixgame2_ode_trajectory = pa.schema([ pa.field("file_name", pa.string()), pa.field("caption", pa.string()), pa.field("media_type", pa.string()), # 'image' or 'video' -]) \ No newline at end of file +]) diff --git a/fastvideo/dataset/latent_datasets.py b/fastvideo/dataset/latent_datasets.py index dc01171d6..277784619 100644 --- a/fastvideo/dataset/latent_datasets.py +++ b/fastvideo/dataset/latent_datasets.py @@ -25,10 +25,8 @@ class LatentDataset(Dataset): self.datase_dir_path = os.path.dirname(json_path) self.video_dir = os.path.join(self.datase_dir_path, "video") self.latent_dir = os.path.join(self.datase_dir_path, "latent") - self.prompt_embed_dir = os.path.join(self.datase_dir_path, - "prompt_embed") - self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path, - "prompt_attention_mask") + self.prompt_embed_dir = os.path.join(self.datase_dir_path, "prompt_embed") + self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path, "prompt_attention_mask") with open(self.json_path) as f: self.data_anno = json.load(f) # json.load(f) already keeps the order @@ -38,15 +36,12 @@ class LatentDataset(Dataset): self.uncond_prompt_embed = torch.zeros(256, 4096).to(torch.float32) self.uncond_prompt_mask = torch.zeros(256).bool() - self.lengths = [ - data_item.get("length", 1) for data_item in self.data_anno - ] + self.lengths = [data_item.get("length", 1) for data_item in self.data_anno] def __getitem__(self, idx): latent_file = self.data_anno[idx]["latent_path"] prompt_embed_file = self.data_anno[idx]["prompt_embed_path"] - prompt_attention_mask_file = self.data_anno[idx][ - "prompt_attention_mask"] + prompt_attention_mask_file = self.data_anno[idx]["prompt_attention_mask"] # load latent = torch.load( os.path.join(self.latent_dir, latent_file), @@ -64,8 +59,7 @@ class LatentDataset(Dataset): weights_only=True, ) prompt_attention_mask = torch.load( - os.path.join(self.prompt_attention_mask_dir, - prompt_attention_mask_file), + os.path.join(self.prompt_attention_mask_dir, prompt_attention_mask_file), map_location="cpu", weights_only=True, ) diff --git a/fastvideo/dataset/ltx2_precomputed_dataset.py b/fastvideo/dataset/ltx2_precomputed_dataset.py index 29c826ee7..ac9c88035 100644 --- a/fastvideo/dataset/ltx2_precomputed_dataset.py +++ b/fastvideo/dataset/ltx2_precomputed_dataset.py @@ -53,39 +53,32 @@ class LTX2PrecomputedDataset(Dataset): def _setup_data_root(data_root: str) -> Path: data_root_path = Path(data_root).expanduser().resolve() if not data_root_path.exists(): - raise FileNotFoundError( - f"Data root directory does not exist: {data_root_path}") + raise FileNotFoundError(f"Data root directory does not exist: {data_root_path}") if (data_root_path / PRECOMPUTED_DIR_NAME).exists(): data_root_path = data_root_path / PRECOMPUTED_DIR_NAME return data_root_path @staticmethod - def _normalize_data_sources( - data_sources: dict[str, str] | list[str] | None, - ) -> dict[str, str]: + def _normalize_data_sources(data_sources: dict[str, str] | list[str] | None, ) -> dict[str, str]: if data_sources is None: return {"latents": "latents", "conditions": "conditions"} if isinstance(data_sources, list): return {source: source for source in data_sources} if isinstance(data_sources, dict): return data_sources.copy() - raise TypeError( - f"data_sources must be dict, list, or None, got {type(data_sources)}") + raise TypeError(f"data_sources must be dict, list, or None, got {type(data_sources)}") def _setup_source_paths(self) -> dict[str, Path]: source_paths: dict[str, Path] = {} for dir_name in self.data_sources: source_path = self.data_root / dir_name if not source_path.exists(): - raise FileNotFoundError( - f"Required {dir_name} directory does not exist: {source_path}") + raise FileNotFoundError(f"Required {dir_name} directory does not exist: {source_path}") source_paths[dir_name] = source_path return source_paths def _discover_samples(self) -> dict[str, list[Path]]: - data_key = ("latents" - if "latents" in self.data_sources else next(iter( - self.data_sources.keys()))) + data_key = ("latents" if "latents" in self.data_sources else next(iter(self.data_sources.keys()))) data_path = self.source_paths[data_key] data_files = list(data_path.glob("**/*.pt")) if not data_files: @@ -100,8 +93,7 @@ class LTX2PrecomputedDataset(Dataset): def _all_source_files_exist(self, data_file: Path, rel_path: Path) -> bool: for dir_name in self.data_sources: - expected_path = self._get_expected_file_path(dir_name, data_file, - rel_path) + expected_path = self._get_expected_file_path(dir_name, data_file, rel_path) if not expected_path.exists(): logger.warning( "No matching %s file found for: %s (expected in: %s)", @@ -112,33 +104,23 @@ class LTX2PrecomputedDataset(Dataset): return False return True - def _get_expected_file_path(self, dir_name: str, data_file: Path, - rel_path: Path) -> Path: + def _get_expected_file_path(self, dir_name: str, data_file: Path, rel_path: Path) -> Path: source_path = self.source_paths[dir_name] if dir_name == "conditions" and data_file.name.startswith("latent_"): return source_path / f"condition_{data_file.stem[7:]}.pt" return source_path / rel_path - def _fill_sample_data_files(self, data_file: Path, rel_path: Path, - sample_files: dict[str, list[Path]]) -> None: + def _fill_sample_data_files(self, data_file: Path, rel_path: Path, sample_files: dict[str, list[Path]]) -> None: for dir_name, output_key in self.data_sources.items(): - expected_path = self._get_expected_file_path(dir_name, data_file, - rel_path) - sample_files[output_key].append( - expected_path.relative_to(self.source_paths[dir_name])) + expected_path = self._get_expected_file_path(dir_name, data_file, rel_path) + sample_files[output_key].append(expected_path.relative_to(self.source_paths[dir_name])) def _validate_setup(self) -> None: if not self.sample_files: - raise ValueError( - "No valid samples found - all data sources must have matching files" - ) - sample_counts = { - key: len(files) - for key, files in self.sample_files.items() - } + raise ValueError("No valid samples found - all data sources must have matching files") + sample_counts = {key: len(files) for key, files in self.sample_files.items()} if len(set(sample_counts.values())) > 1: - raise ValueError( - f"Mismatched sample counts across sources: {sample_counts}") + raise ValueError(f"Mismatched sample counts across sources: {sample_counts}") def __len__(self) -> int: first_key = next(iter(self.sample_files.keys())) @@ -156,8 +138,7 @@ class LTX2PrecomputedDataset(Dataset): data = self._normalize_video_latents(data) result[output_key] = data except Exception as e: - raise RuntimeError( - f"Failed to load {output_key} from {file_path}: {e}") from e + raise RuntimeError(f"Failed to load {output_key} from {file_path}: {e}") from e result["idx"] = index return result diff --git a/fastvideo/dataset/parquet_dataset_iterable_style.py b/fastvideo/dataset/parquet_dataset_iterable_style.py index a4fbf1409..7c9b80b72 100644 --- a/fastvideo/dataset/parquet_dataset_iterable_style.py +++ b/fastvideo/dataset/parquet_dataset_iterable_style.py @@ -10,8 +10,7 @@ from torch.utils.data import IterableDataset, get_worker_info from torchdata.stateful_dataloader import StatefulDataLoader from fastvideo.dataset.utils import collate_latents_embs_masks -from fastvideo.distributed import (get_sp_world_size, get_world_group, - get_world_rank, get_world_size) +from fastvideo.distributed import (get_sp_world_size, get_world_group, get_world_rank, get_world_size) from fastvideo.logger import init_logger logger = init_logger(__name__) @@ -19,8 +18,7 @@ logger = init_logger(__name__) class BatchIterator: # TODO: Implement state_dict and load_state_dict to support resume. - def __init__(self, files, batch_size, text_padding_length, keys, - worker_num_samples, read_batch_size): + def __init__(self, files, batch_size, text_padding_length, keys, worker_num_samples, read_batch_size): self.files = files self.batch_size = batch_size self.text_padding_length = text_padding_length @@ -90,25 +88,18 @@ class LatentsParquetIterStyleDataset(IterableDataset): self.path, self.num_sp_groups, num_workers, seed) if drop_last: - self.worker_num_samples = min( - shard_total_samples) // batch_size * batch_size + self.worker_num_samples = min(shard_total_samples) // batch_size * batch_size # Assign files to current rank's SP group ith_sp_group = self.global_rank // self.sp_world_size - self.sp_group_parquet_files = shard_parquet_files[ith_sp_group::self - .num_sp_groups] - self.sp_group_parquet_lengths = shard_parquet_lengths[ - ith_sp_group::self.num_sp_groups] - self.sp_group_num_samples = shard_total_samples[ith_sp_group::self. - num_sp_groups] - logger.info( - "In total %d parquet files, %d samples, after sharding we retain %d samples due to drop_last", - sum([len(shard) for shard in shard_parquet_files]), - sum(shard_total_samples), - self.worker_num_samples * self.num_sp_groups * num_workers) + self.sp_group_parquet_files = shard_parquet_files[ith_sp_group::self.num_sp_groups] + self.sp_group_parquet_lengths = shard_parquet_lengths[ith_sp_group::self.num_sp_groups] + self.sp_group_num_samples = shard_total_samples[ith_sp_group::self.num_sp_groups] + logger.info("In total %d parquet files, %d samples, after sharding we retain %d samples due to drop_last", + sum([len(shard) for shard in shard_parquet_files]), sum(shard_total_samples), + self.worker_num_samples * self.num_sp_groups * num_workers) else: raise ValueError("drop_last must be True") - logger.info("Each dataloader worker will load %d samples", - self.worker_num_samples) + logger.info("Each dataloader worker will load %d samples", self.worker_num_samples) def __iter__(self): worker_info = get_worker_info() @@ -116,20 +107,18 @@ class LatentsParquetIterStyleDataset(IterableDataset): worker_files = self.sp_group_parquet_files[worker_id] - batch_iterator = BatchIterator( - files=worker_files, - batch_size=self.batch_size, - text_padding_length=self.text_padding_length, - keys=self.keys, - worker_num_samples=self.worker_num_samples, - read_batch_size=self.read_batch_size) # type: ignore + batch_iterator = BatchIterator(files=worker_files, + batch_size=self.batch_size, + text_padding_length=self.text_padding_length, + keys=self.keys, + worker_num_samples=self.worker_num_samples, + read_batch_size=self.read_batch_size) # type: ignore yield from batch_iterator if batch_iterator.processed_samples != self.worker_num_samples: - raise ValueError( - "Rank %d, Worker %d: Not enough samples to process, this should not happen", - self.global_rank, worker_id) + raise ValueError("Rank %d, Worker %d: Not enough samples to process, this should not happen", + self.global_rank, worker_id) def shard_parquet_files_across_sp_groups_and_workers( @@ -154,8 +143,7 @@ def shard_parquet_files_across_sp_groups_and_workers( - List of dictionaries mapping file paths to their lengths """ # Check if sharding plan already exists - sharding_info_dir = os.path.join( - path, f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers") + sharding_info_dir = os.path.join(path, f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers") # Only rank 0 handles cache checking and file scanning if get_world_rank() == 0: @@ -168,17 +156,11 @@ def shard_parquet_files_across_sp_groups_and_workers( if os.path.exists(sharding_info_dir): logger.info("Loading sharding plan from %s", sharding_info_dir) try: - with open( - os.path.join(sharding_info_dir, - "shard_parquet_files.pkl"), "rb") as f: + with open(os.path.join(sharding_info_dir, "shard_parquet_files.pkl"), "rb") as f: shard_parquet_files = pickle.load(f) - with open( - os.path.join(sharding_info_dir, - "shard_total_samples.pkl"), "rb") as f: + with open(os.path.join(sharding_info_dir, "shard_total_samples.pkl"), "rb") as f: shard_total_samples = pickle.load(f) - with open( - os.path.join(sharding_info_dir, - "shard_parquet_lengths.pkl"), "rb") as f: + with open(os.path.join(sharding_info_dir, "shard_parquet_lengths.pkl"), "rb") as f: shard_parquet_lengths = pickle.load(f) cache_loaded = True logger.info("Successfully loaded sharding plan") @@ -210,8 +192,7 @@ def shard_parquet_files_across_sp_groups_and_workers( lengths.append(pq.ParquetFile(file).metadata.num_rows) total_samples = sum(lengths) - logger.info("Found %d files with %d total samples", - len(parquet_files), total_samples) + logger.info("Found %d files with %d total samples", len(parquet_files), total_samples) # Sort files by length for better balancing sorted_indices = np.argsort(lengths) @@ -226,9 +207,7 @@ def shard_parquet_files_across_sp_groups_and_workers( # Distribute files to shards using a greedy approach logger.info("Distributing files to shards...") - for file, length in zip(reversed(sorted_files), - reversed(sorted_lengths), - strict=True): + for file, length in zip(reversed(sorted_files), reversed(sorted_lengths), strict=True): # Find shard with minimum current length target_shard = np.argmin(shard_total_samples) shard_parquet_files[target_shard].append(file) @@ -241,17 +220,11 @@ def shard_parquet_files_across_sp_groups_and_workers( # Save the sharding plan os.makedirs(sharding_info_dir, exist_ok=True) - with open( - os.path.join(sharding_info_dir, "shard_parquet_files.pkl"), - "wb") as f: + with open(os.path.join(sharding_info_dir, "shard_parquet_files.pkl"), "wb") as f: pickle.dump(shard_parquet_files, f) - with open( - os.path.join(sharding_info_dir, "shard_total_samples.pkl"), - "wb") as f: + with open(os.path.join(sharding_info_dir, "shard_total_samples.pkl"), "wb") as f: pickle.dump(shard_total_samples, f) - with open( - os.path.join(sharding_info_dir, - "shard_parquet_lengths.pkl"), "wb") as f: + with open(os.path.join(sharding_info_dir, "shard_parquet_lengths.pkl"), "wb") as f: pickle.dump(shard_parquet_lengths, f) logger.info("Saved sharding info to %s", sharding_info_dir) @@ -260,41 +233,35 @@ def shard_parquet_files_across_sp_groups_and_workers( world_group.barrier() # Now all ranks load the sharding plan (it should exist and be valid now) - logger.info("Loading sharding plan from %s after barrier", - sharding_info_dir) - with open(os.path.join(sharding_info_dir, "shard_parquet_files.pkl"), - "rb") as f: + logger.info("Loading sharding plan from %s after barrier", sharding_info_dir) + with open(os.path.join(sharding_info_dir, "shard_parquet_files.pkl"), "rb") as f: shard_parquet_files = pickle.load(f) - with open(os.path.join(sharding_info_dir, "shard_total_samples.pkl"), - "rb") as f: + with open(os.path.join(sharding_info_dir, "shard_total_samples.pkl"), "rb") as f: shard_total_samples = pickle.load(f) - with open(os.path.join(sharding_info_dir, "shard_parquet_lengths.pkl"), - "rb") as f: + with open(os.path.join(sharding_info_dir, "shard_parquet_lengths.pkl"), "rb") as f: shard_parquet_lengths = pickle.load(f) return shard_parquet_files, shard_total_samples, shard_parquet_lengths def build_parquet_iterable_style_dataloader( - path: str, - batch_size: int, - num_data_workers: int, - cfg_rate: float = 0.0, - drop_last: bool = True, - text_padding_length: int = 512, - seed: int = 42, - read_batch_size: int = 32 -) -> tuple[LatentsParquetIterStyleDataset, StatefulDataLoader]: + path: str, + batch_size: int, + num_data_workers: int, + cfg_rate: float = 0.0, + drop_last: bool = True, + text_padding_length: int = 512, + seed: int = 42, + read_batch_size: int = 32) -> tuple[LatentsParquetIterStyleDataset, StatefulDataLoader]: """Build a dataloader for the LatentsParquetIterStyleDataset.""" - dataset = LatentsParquetIterStyleDataset( - path=path, - batch_size=batch_size, - cfg_rate=cfg_rate, - num_workers=num_data_workers, - drop_last=drop_last, - text_padding_length=text_padding_length, - seed=seed, - read_batch_size=read_batch_size) + dataset = LatentsParquetIterStyleDataset(path=path, + batch_size=batch_size, + cfg_rate=cfg_rate, + num_workers=num_data_workers, + drop_last=drop_last, + text_padding_length=text_padding_length, + seed=seed, + read_batch_size=read_batch_size) loader = StatefulDataLoader( dataset, diff --git a/fastvideo/dataset/parquet_dataset_map_style.py b/fastvideo/dataset/parquet_dataset_map_style.py index be41ab46e..f99937f75 100644 --- a/fastvideo/dataset/parquet_dataset_map_style.py +++ b/fastvideo/dataset/parquet_dataset_map_style.py @@ -16,8 +16,7 @@ from torchdata.stateful_dataloader import StatefulDataLoader from fastvideo.platforms import current_platform from fastvideo.dataset.utils import collate_rows_from_parquet_schema -from fastvideo.distributed import (get_sp_world_size, get_world_group, - get_world_rank, get_world_size) +from fastvideo.distributed import (get_sp_world_size, get_world_group, get_world_rank, get_world_size) from fastvideo.logger import init_logger logger = init_logger(__name__) @@ -64,26 +63,20 @@ class DP_SP_BatchSampler(Sampler[list[int]]): # 3. Prevents uneven batch sizes across SP groups at end of epoch num_batches = self.dataset_size // self.batch_size num_global_batches = num_batches // self.num_sp_groups - global_indices = global_indices[:num_global_batches * - self.num_sp_groups * - self.batch_size] + global_indices = global_indices[:num_global_batches * self.num_sp_groups * self.batch_size] else: if self.dataset_size % (self.num_sp_groups * self.batch_size) != 0: # add more indices to make it divisible by (batch_size * num_sp_groups) - padding_size = self.num_sp_groups * self.batch_size - ( - self.dataset_size % (self.num_sp_groups * self.batch_size)) - logger.info("Padding the dataset from %d to %d", - self.dataset_size, self.dataset_size + padding_size) - global_indices = torch.cat( - [global_indices, global_indices[:padding_size]]) + padding_size = self.num_sp_groups * self.batch_size - (self.dataset_size % + (self.num_sp_groups * self.batch_size)) + logger.info("Padding the dataset from %d to %d", self.dataset_size, self.dataset_size + padding_size) + global_indices = torch.cat([global_indices, global_indices[:padding_size]]) # shard the indices to each sp group ith_sp_group = self.global_rank // self.sp_world_size - sp_group_local_indices = global_indices[ith_sp_group::self. - num_sp_groups] + sp_group_local_indices = global_indices[ith_sp_group::self.num_sp_groups] self.sp_group_local_indices = sp_group_local_indices - logger.info("Dataset size for each sp group: %d", - len(sp_group_local_indices)) + logger.info("Dataset size for each sp group: %d", len(sp_group_local_indices)) def __iter__(self): indices = self.sp_group_local_indices @@ -98,11 +91,7 @@ class DP_SP_BatchSampler(Sampler[list[int]]): def _parse_data_path_specs(path: str | Sequence[str] | dict[str, int]) -> list[tuple[str, int]]: """Parse one or more dataset roots with old-framework repeat counts.""" if isinstance(path, dict): - return [ - (str(root), int(repeat)) - for root, repeat in path.items() - if int(repeat) > 0 - ] + return [(str(root), int(repeat)) for root, repeat in path.items() if int(repeat) > 0] if isinstance(path, Sequence) and not isinstance(path, str): return [(str(root), 1) for root in path] @@ -132,11 +121,9 @@ def get_parquet_files_and_length(path: str | Sequence[str] | dict[str, int]): all_file_names.extend(file_names) all_lengths.extend(lengths) if not all_file_names: - raise FileNotFoundError( - "No parquet files found under dataset paths: " - f"{path}. " - "Please verify these paths point to preprocessed parquet data." - ) + raise FileNotFoundError("No parquet files found under dataset paths: " + f"{path}. " + "Please verify these paths point to preprocessed parquet data.") file_lengths = sorted( zip(all_file_names, all_lengths, strict=True), key=lambda x: x[0], @@ -162,19 +149,13 @@ def get_parquet_files_and_length(path: str | Sequence[str] | dict[str, int]): with open(cache_file, "rb") as f: file_names_sorted, lengths_sorted = pickle.load(f) file_names_sorted = tuple( - os.path.realpath( - os.path.join(os.getcwd(), p) - if not os.path.isabs(p) else p) + os.path.realpath(os.path.join(os.getcwd(), p) if not os.path.isabs(p) else p) for p in file_names_sorted) files_outside_dataset_root = [ file_path for file_path in file_names_sorted - if os.path.commonpath([dataset_root, file_path - ]) != dataset_root - ] - missing_files = [ - file_path for file_path in file_names_sorted - if not os.path.exists(file_path) + if os.path.commonpath([dataset_root, file_path]) != dataset_root ] + missing_files = [file_path for file_path in file_names_sorted if not os.path.exists(file_path)] if files_outside_dataset_root: logger.warning( "Cached parquet file list points outside dataset root " @@ -210,20 +191,15 @@ def get_parquet_files_and_length(path: str | Sequence[str] | dict[str, int]): file_path = os.path.realpath(os.path.join(root, file)) file_names.append(file_path) if len(file_names) == 0: - raise FileNotFoundError( - "No parquet files found under dataset path: " - f"{path}. " - "Please verify this path points to preprocessed parquet " - "data.") - for file_path in tqdm.tqdm( - file_names, desc="Reading parquet files to get lengths"): + raise FileNotFoundError("No parquet files found under dataset path: " + f"{path}. " + "Please verify this path points to preprocessed parquet " + "data.") + for file_path in tqdm.tqdm(file_names, desc="Reading parquet files to get lengths"): num_rows = pq.ParquetFile(file_path).metadata.num_rows lengths.append(num_rows) # sort according to file name to ensure all rank has the same order - file_names_sorted, lengths_sorted = zip(*sorted(zip(file_names, - lengths, - strict=True), - key=lambda x: x[0]), + file_names_sorted, lengths_sorted = zip(*sorted(zip(file_names, lengths, strict=True), key=lambda x: x[0]), strict=True) # Save the cache os.makedirs(cache_dir, exist_ok=True) @@ -240,20 +216,17 @@ def get_parquet_files_and_length(path: str | Sequence[str] | dict[str, int]): with open(cache_file, "rb") as f: file_names_sorted, lengths_sorted = pickle.load(f) if len(file_names_sorted) == 0: - raise RuntimeError( - "Cached parquet metadata is empty after synchronization at " - f"{cache_file}. " - "Please verify the dataset path and regenerate cache.") + raise RuntimeError("Cached parquet metadata is empty after synchronization at " + f"{cache_file}. " + "Please verify the dataset path and regenerate cache.") if len(file_names_sorted) != len(lengths_sorted): - raise RuntimeError( - "Cached parquet metadata is corrupted at " - f"{cache_file}: file count and length count do not match.") + raise RuntimeError("Cached parquet metadata is corrupted at " + f"{cache_file}: file count and length count do not match.") return file_names_sorted, lengths_sorted -def read_row_from_parquet_file(parquet_files: list[str], global_row_idx: int, - lengths: list[int]) -> dict[str, Any]: +def read_row_from_parquet_file(parquet_files: list[str], global_row_idx: int, lengths: list[int]) -> dict[str, Any]: ''' Read a row from a parquet file. Args: @@ -274,8 +247,7 @@ def read_row_from_parquet_file(parquet_files: list[str], global_row_idx: int, cumulative += lengths[file_index] else: # If we reach here, global_row_idx is out of bounds - raise IndexError( - f"global_row_idx {global_row_idx} is out of bounds for dataset") + raise IndexError(f"global_row_idx {global_row_idx} is out of bounds for dataset") parquet_file = pq.ParquetFile(parquet_files[file_index]) @@ -294,9 +266,7 @@ def read_row_from_parquet_file(parquet_files: list[str], global_row_idx: int, cumulative += num_rows else: # If we reach here, local_row_idx is out of bounds for this parquet file - raise IndexError( - f"local_row_idx {local_row_idx} is out of bounds for parquet file {parquet_files[file_index]}" - ) + raise IndexError(f"local_row_idx {local_row_idx} is out of bounds for parquet file {parquet_files[file_index]}") row_group = parquet_file.read_row_group(row_group_index).to_pydict() row_dict = {k: v[local_index] for k, v in row_group.items()} @@ -333,8 +303,7 @@ class LatentsParquetMapStyleDataset(Dataset): self.seed = seed # Create a seeded random generator for deterministic CFG self.rng = random.Random(seed) - logger.info("Initializing LatentsParquetMapStyleDataset with path: %s", - path) + logger.info("Initializing LatentsParquetMapStyleDataset with path: %s", path) self.parquet_files, self.lengths = get_parquet_files_and_length(path) self.batch = batch_size self.text_padding_length = text_padding_length @@ -348,11 +317,9 @@ class LatentsParquetMapStyleDataset(Dataset): drop_first_row=drop_first_row, seed=seed, ) - logger.info("Dataset initialized with %d parquet files and %d rows", - len(self.parquet_files), sum(self.lengths)) + logger.info("Dataset initialized with %d parquet files and %d rows", len(self.parquet_files), sum(self.lengths)) - def get_validation_negative_prompt( - self) -> tuple[torch.Tensor, torch.Tensor, str]: + def get_validation_negative_prompt(self) -> tuple[torch.Tensor, torch.Tensor, str]: """ Get the negative prompt for validation. This method ensures the negative prompt is loaded and cached properly. @@ -363,8 +330,7 @@ class LatentsParquetMapStyleDataset(Dataset): file_path = self.parquet_files[0] row_idx = 0 # Read the negative prompt data - row_dict = read_row_from_parquet_file([file_path], row_idx, - [self.lengths[0]]) + row_dict = read_row_from_parquet_file([file_path], row_idx, [self.lengths[0]]) batch = collate_rows_from_parquet_schema([row_dict], self.parquet_schema, @@ -377,8 +343,7 @@ class LatentsParquetMapStyleDataset(Dataset): if len(negative_prompt_embedding.shape) == 2: negative_prompt_embedding = negative_prompt_embedding.unsqueeze(0) if len(negative_prompt_attention_mask.shape) == 1: - negative_prompt_attention_mask = negative_prompt_attention_mask.unsqueeze( - 0).unsqueeze(0) + negative_prompt_attention_mask = negative_prompt_attention_mask.unsqueeze(0).unsqueeze(0) return negative_prompt_embedding, negative_prompt_attention_mask, negative_prompt @@ -387,10 +352,7 @@ class LatentsParquetMapStyleDataset(Dataset): """ Batch fetch using read_row_from_parquet_file for each index. """ - rows = [ - read_row_from_parquet_file(self.parquet_files, idx, self.lengths) - for idx in indices - ] + rows = [read_row_from_parquet_file(self.parquet_files, idx, self.lengths) for idx in indices] # Inject sample indices for deterministic CFG dropout # that is reproducible across checkpoint resume. @@ -407,6 +369,7 @@ class LatentsParquetMapStyleDataset(Dataset): def __len__(self): return sum(self.lengths) + # ──────────────────────────────────────────────────────────────────────────── # 3. Loader helper – everything else stays just like your original trainer # ──────────────────────────────────────────────────────────────────────────── @@ -414,25 +377,23 @@ def passthrough(batch): return batch -def build_parquet_map_style_dataloader( - path, - batch_size, - num_data_workers, - parquet_schema, - cfg_rate=0.0, - drop_last=True, - drop_first_row=False, - text_padding_length=512, - seed=42) -> tuple[LatentsParquetMapStyleDataset, StatefulDataLoader]: - dataset = LatentsParquetMapStyleDataset( - path, - batch_size, - cfg_rate=cfg_rate, - drop_last=drop_last, - drop_first_row=drop_first_row, - text_padding_length=text_padding_length, - parquet_schema=parquet_schema, - seed=seed) +def build_parquet_map_style_dataloader(path, + batch_size, + num_data_workers, + parquet_schema, + cfg_rate=0.0, + drop_last=True, + drop_first_row=False, + text_padding_length=512, + seed=42) -> tuple[LatentsParquetMapStyleDataset, StatefulDataLoader]: + dataset = LatentsParquetMapStyleDataset(path, + batch_size, + cfg_rate=cfg_rate, + drop_last=drop_last, + drop_first_row=drop_first_row, + text_padding_length=text_padding_length, + parquet_schema=parquet_schema, + seed=seed) loader = StatefulDataLoader( dataset, diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py index 9d1ec44e7..85696e53a 100644 --- a/fastvideo/dataset/preprocessing_datasets.py +++ b/fastvideo/dataset/preprocessing_datasets.py @@ -185,24 +185,19 @@ class FrameSamplingStage(DatasetFilterStage): num_frames = math.ceil(batch.fps * batch.duration) # Check if video is too long - if (num_frames / batch.fps > self.video_length_tolerance_range * - (self.num_frames / self.train_fps * self.speed_factor)): + if (num_frames / batch.fps + > self.video_length_tolerance_range * (self.num_frames / self.train_fps * self.speed_factor)): return False # Resample frame indices to check length frame_interval = batch.fps / self.train_fps start_frame_idx = 0 - frame_indices = np.arange(start_frame_idx, num_frames, - frame_interval).astype(int) + frame_indices = np.arange(start_frame_idx, num_frames, frame_interval).astype(int) # Filter short videos - return not (len(frame_indices) < self.num_frames - and self.rng.random() < self.drop_short_ratio) + return not (len(frame_indices) < self.num_frames and self.rng.random() < self.drop_short_ratio) - def process(self, - batch: PreprocessBatch, - temporal_sample_fn=None, - **kwargs) -> PreprocessBatch: + def process(self, batch: PreprocessBatch, temporal_sample_fn=None, **kwargs) -> PreprocessBatch: """ Process frame sampling for video data items. @@ -225,8 +220,7 @@ class FrameSamplingStage(DatasetFilterStage): # Resample frame indices frame_interval = batch.fps / self.train_fps start_frame_idx = 0 - frame_indices = np.arange(start_frame_idx, batch.num_frames, - frame_interval).astype(int) + frame_indices = np.arange(start_frame_idx, batch.num_frames, frame_interval).astype(int) # Temporal crop if too long if len(frame_indices) > self.num_frames: @@ -264,8 +258,7 @@ class VideoTransformStage(DatasetStage): assert os.path.exists(batch.path), f"file {batch.path} do not exist!" assert batch.sample_frame_index is not None, "Frame indices must be set before transformation" - torchvision_video, _, metadata = torchvision.io.read_video( - batch.path, output_format="TCHW") + torchvision_video, _, metadata = torchvision.io.read_video(batch.path, output_format="TCHW") video = torchvision_video[batch.sample_frame_index] if self.transform is not None: video = self.transform(video) @@ -317,11 +310,7 @@ class ImageTransformStage(DatasetStage): class TextEncodingStage(DatasetStage): """Stage for text tokenization and encoding.""" - def __init__(self, - tokenizer, - text_max_length: int, - cfg_rate: float = 0.0, - seed: int = 42): + def __init__(self, tokenizer, text_max_length: int, cfg_rate: float = 0.0, seed: int = 42): self.tokenizer = tokenizer self.text_max_length = text_max_length self.cfg_rate = cfg_rate @@ -360,8 +349,7 @@ class TextEncodingStage(DatasetStage): return batch -class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, - torch.distributed.checkpoint.stateful.Stateful): +class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, torch.distributed.checkpoint.stateful.Stateful): """ Merged dataset for video and caption data with stage-based processing. Assumes that data_merge_path is a txt file with the following format: @@ -415,8 +403,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, tokenizer = None if os.path.exists(tokenizer_path): try: - tokenizer = AutoTokenizer.from_pretrained( - tokenizer_path, cache_dir=args.cache_dir) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, cache_dir=args.cache_dir) except (ValueError, OSError): pass @@ -426,26 +413,22 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, # Process metadata self.processed_batches = self._process_metadata() - def _init_stages(self, args, transform, transform_topcrop, - tokenizer) -> None: + def _init_stages(self, args, transform, transform_topcrop, tokenizer) -> None: """Initialize all processing stages.""" self.validation_stage = DataValidationStage() - self.frame_sampling_stage = FrameSamplingStage( - num_frames=args.num_frames, - train_fps=args.train_fps, - speed_factor=args.speed_factor, - video_length_tolerance_range=args.video_length_tolerance_range, - drop_short_ratio=args.drop_short_ratio, - seed=self.seed) + self.frame_sampling_stage = FrameSamplingStage(num_frames=args.num_frames, + train_fps=args.train_fps, + speed_factor=args.speed_factor, + video_length_tolerance_range=args.video_length_tolerance_range, + drop_short_ratio=args.drop_short_ratio, + seed=self.seed) self.video_transform_stage = VideoTransformStage(transform) - self.image_transform_stage = ImageTransformStage( - transform, transform_topcrop) + self.image_transform_stage = ImageTransformStage(transform, transform_topcrop) if tokenizer is not None: - self.text_encoding_stage = TextEncodingStage( - tokenizer=tokenizer, - text_max_length=args.text_max_length, - cfg_rate=args.training_cfg_rate, - seed=self.seed) + self.text_encoding_stage = TextEncodingStage(tokenizer=tokenizer, + text_max_length=args.text_max_length, + cfg_rate=args.training_cfg_rate, + seed=self.seed) else: self.text_encoding_stage = None @@ -453,13 +436,9 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, """Load raw data from JSON files.""" # Read folder-annotation pairs with open(self.data_merge_path) as f: - folder_anno_pairs = [ - line.strip().split(",") for line in f if line.strip() - ] - assert len( - folder_anno_pairs) == 1, "Only support one folder-annotation pair" - assert len(folder_anno_pairs[0] - ) == 2, "Folder-annotation pair should have two elements" + folder_anno_pairs = [line.strip().split(",") for line in f if line.strip()] + assert len(folder_anno_pairs) == 1, "Only support one folder-annotation pair" + assert len(folder_anno_pairs[0]) == 2, "Folder-annotation pair should have two elements" folder, annotation_file = folder_anno_pairs[0] data_items: list[dict] = [] @@ -480,10 +459,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, processed_batches = [] # Initialize counters - filter_counts = { - "validation_failed": 0, - "frame_sampling_failed": 0 - } + filter_counts = {"validation_failed": 0, "frame_sampling_failed": 0} sample_num_frames: list[int] = [] for item in raw_data: @@ -499,19 +475,16 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, continue # Apply frame sampling processing - batch = self.frame_sampling_stage.process( - batch, temporal_sample_fn=self.temporal_sample) + batch = self.frame_sampling_stage.process(batch, temporal_sample_fn=self.temporal_sample) processed_batches.append(batch) assert batch.sample_num_frames is not None sample_num_frames.append(batch.sample_num_frames) - self._log_filtering_stats(filter_counts, sample_num_frames, - len(raw_data), len(processed_batches)) + self._log_filtering_stats(filter_counts, sample_num_frames, len(raw_data), len(processed_batches)) return processed_batches - def _apply_filter_stages(self, batch: PreprocessBatch, - filter_counts: dict[str, int]) -> bool: + def _apply_filter_stages(self, batch: PreprocessBatch, filter_counts: dict[str, int]) -> bool: """Apply all filter stages and update counters. Returns True if batch should be kept.""" if not self.validation_stage.should_keep(batch): filter_counts["validation_failed"] += 1 @@ -523,16 +496,13 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, return True - def _log_filtering_stats(self, filter_counts: dict[str, int], - sample_num_frames: list[int], before_count: int, + def _log_filtering_stats(self, filter_counts: dict[str, int], sample_num_frames: list[int], before_count: int, after_count: int): """Log filtering statistics.""" logger.info( "validation_failed: %d, frame_sampling_failed: %d, " - "Counter(sample_num_frames): %s, before filter: %d, after filter: %d", - filter_counts['validation_failed'], - filter_counts['frame_sampling_failed'], Counter(sample_num_frames), - before_count, after_count) + "Counter(sample_num_frames): %s, before filter: %d, after filter: %d", filter_counts['validation_failed'], + filter_counts['frame_sampling_failed'], Counter(sample_num_frames), before_count, after_count) def __iter__(self): """Iterate through processed data items.""" @@ -566,7 +536,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, # Add video-specific fields if batch.is_video: result.update({"fps": batch.fps, "duration": batch.duration}) - + # Add action_path if batch.action_path: result["action_path"] = batch.action_path @@ -582,8 +552,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, self.processed_batches = state_dict["processed_batches"] -class TextDataset(torch.utils.data.IterableDataset, - torch.distributed.checkpoint.stateful.Stateful): +class TextDataset(torch.utils.data.IterableDataset, torch.distributed.checkpoint.stateful.Stateful): """ Text-only dataset for processing prompts from a simple text file. @@ -596,11 +565,7 @@ class TextDataset(torch.utils.data.IterableDataset, This dataset processes text data through text encoding stages only. """ - def __init__(self, - data_merge_path: str, - args, - start_idx: int = 0, - seed: int = 42): + def __init__(self, data_merge_path: str, args, start_idx: int = 0, seed: int = 42): self.data_merge_path = data_merge_path self.start_idx = start_idx self.args = args @@ -608,15 +573,13 @@ class TextDataset(torch.utils.data.IterableDataset, # Initialize tokenizer tokenizer_path = os.path.join(args.model_path, "tokenizer") - tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, - cache_dir=args.cache_dir) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, cache_dir=args.cache_dir) # Initialize text encoding stage - self.text_encoding_stage = TextEncodingStage( - tokenizer=tokenizer, - text_max_length=args.text_max_length, - cfg_rate=getattr(args, 'training_cfg_rate', 0.0), - seed=self.seed) + self.text_encoding_stage = TextEncodingStage(tokenizer=tokenizer, + text_max_length=args.text_max_length, + cfg_rate=getattr(args, 'training_cfg_rate', 0.0), + seed=self.seed) # Process text data self.processed_batches = self._process_text_data() @@ -629,7 +592,7 @@ class TextDataset(torch.utils.data.IterableDataset, line = line.strip() if line: # Skip empty lines prompts.append(line) - + logger.info(f"Loaded {len(prompts)} text prompts from {self.data_merge_path}") return prompts @@ -648,9 +611,8 @@ class TextDataset(torch.utils.data.IterableDataset, duration=None, num_frames=0, sample_frame_index=None, - sample_num_frames=0 - ) - + sample_num_frames=0) + processed_batches.append(batch) logger.info(f"Processed {len(processed_batches)} text batches") diff --git a/fastvideo/dataset/transform.py b/fastvideo/dataset/transform.py index 53c9c49f3..c9a215e8d 100644 --- a/fastvideo/dataset/transform.py +++ b/fastvideo/dataset/transform.py @@ -26,9 +26,7 @@ def crop(clip, i, j, h, w) -> torch.Tensor: def resize(clip, target_size, interpolation_mode) -> torch.Tensor: if len(target_size) != 2: - raise ValueError( - f"target size should be tuple (height, width), instead got {target_size}" - ) + raise ValueError(f"target size should be tuple (height, width), instead got {target_size}") return torch.nn.functional.interpolate( clip, size=target_size, @@ -68,8 +66,7 @@ def normalize_video(clip) -> torch.Tensor: """ _is_tensor_video_clip(clip) if not clip.dtype == torch.uint8: - raise TypeError( - f"clip tensor should have data type uint8. Got {clip.dtype}") + raise TypeError(f"clip tensor should have data type uint8. Got {clip.dtype}") # return clip.float().permute(3, 0, 1, 2) / 255.0 return clip.float() / 255.0 @@ -87,8 +84,7 @@ class CenterCropResizeVideo: interpolation_mode="bilinear", ) -> None: if len(size) != 2: - raise ValueError( - f"size should be tuple (height, width), instead got {size}") + raise ValueError(f"size should be tuple (height, width), instead got {size}") self.size = size self.top_crop = top_crop self.interpolation_mode = interpolation_mode @@ -101,10 +97,7 @@ class CenterCropResizeVideo: torch.tensor: scale resized / center cropped video clip. size is (T, C, crop_size, crop_size) """ - clip_center_crop = center_crop_th_tw(clip, - self.size[0], - self.size[1], - top_crop=self.top_crop) + clip_center_crop = center_crop_th_tw(clip, self.size[0], self.size[1], top_crop=self.top_crop) clip_center_crop_resize = resize( clip_center_crop, target_size=self.size, diff --git a/fastvideo/dataset/utils.py b/fastvideo/dataset/utils.py index 9e5315ae5..efe8d8e2b 100644 --- a/fastvideo/dataset/utils.py +++ b/fastvideo/dataset/utils.py @@ -15,16 +15,12 @@ def pad(t: torch.Tensor, padding_length: int) -> tuple[torch.Tensor, torch.Tenso L, D = t.shape if padding_length > L: # pad pad = torch.zeros(padding_length - L, D, dtype=t.dtype, device=t.device) - return torch.cat([t, pad], 0), torch.cat( - [torch.ones(L), torch.zeros(padding_length - L)], 0) + return torch.cat([t, pad], 0), torch.cat([torch.ones(L), torch.zeros(padding_length - L)], 0) else: # crop return t[:padding_length], torch.ones(padding_length) -def get_torch_tensors_from_row_dict(row_dict, - keys, - cfg_rate, - rng=None) -> dict[str, Any]: +def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate, rng=None) -> dict[str, Any]: """ Get the latents and prompts from a row dictionary. """ @@ -46,8 +42,7 @@ def get_torch_tensors_from_row_dict(row_dict, bytes = row_dict[f"{key}_bytes"] # TODO (peiyuan): read precision - if key == 'text_embedding' and (rng.random() - if rng else random.random()) < cfg_rate: + if key == 'text_embedding' and (rng.random() if rng else random.random()) < cfg_rate: data = np.zeros(shape, dtype=np.float32) else: data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy() @@ -60,12 +55,11 @@ def get_torch_tensors_from_row_dict(row_dict, return return_dict -def collate_latents_embs_masks( - batch_to_process, - text_padding_length, - keys, - cfg_rate=0.0, - rng=None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[str]]: +def collate_latents_embs_masks(batch_to_process, + text_padding_length, + keys, + cfg_rate=0.0, + rng=None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[str]]: # Initialize tensors to hold padded embeddings and masks all_latents = [] all_embs = [] @@ -153,19 +147,12 @@ def collate_rows_from_parquet_schema(rows, # Deterministic per-sample CFG dropout # using sample index (resume-safe). drop = False - if (tensor_name == 'text_embedding' - and cfg_rate > 0): - sample_idx = row.get( - "_sample_index") + if (tensor_name == 'text_embedding' and cfg_rate > 0): + sample_idx = row.get("_sample_index") if sample_idx is not None: - drop = (random.Random( - seed ^ sample_idx - ).random() < cfg_rate) + drop = (random.Random(seed ^ sample_idx).random() < cfg_rate) else: - drop = ((rng.random() - if rng else - random.random()) - < cfg_rate) + drop = ((rng.random() if rng else random.random()) < cfg_rate) if drop: data = np.zeros(shape, dtype=np.float32) else: @@ -197,10 +184,7 @@ def collate_rows_from_parquet_schema(rows, attention_masks.append(mask) else: # Handle empty embeddings - assume default embedding dimension - padded_tensors.append( - torch.zeros(text_padding_length, - 768, - dtype=torch.bfloat16)) + padded_tensors.append(torch.zeros(text_padding_length, 768, dtype=torch.bfloat16)) attention_masks.append(torch.zeros(text_padding_length)) batch_data[tensor_name] = torch.stack(padded_tensors) @@ -211,16 +195,11 @@ def collate_rows_from_parquet_schema(rows, try: batch_data[tensor_name] = torch.stack(tensor_list) except ValueError as e: - shapes = [ - t.shape - if t is not None and hasattr(t, 'shape') else 'None/Invalid' - for t in tensor_list - ] - raise ValueError( - f"Failed to stack tensors for field '{tensor_name}'. " - f"Tensor shapes: {shapes}. " - f"All tensors in a batch must have compatible shapes. " - f"Original error: {e}") from e + shapes = [t.shape if t is not None and hasattr(t, 'shape') else 'None/Invalid' for t in tensor_list] + raise ValueError(f"Failed to stack tensors for field '{tensor_name}'. " + f"Tensor shapes: {shapes}. " + f"All tensors in a batch must have compatible shapes. " + f"Original error: {e}") from e # Process metadata fields into info_list info_list = [] diff --git a/fastvideo/dataset/validation_dataset.py b/fastvideo/dataset/validation_dataset.py index 324e04f35..755df34f6 100644 --- a/fastvideo/dataset/validation_dataset.py +++ b/fastvideo/dataset/validation_dataset.py @@ -7,8 +7,7 @@ import datasets import numpy as np from torch.utils.data import IterableDataset -from fastvideo.distributed import (get_sp_world_size, get_world_rank, - get_world_size) +from fastvideo.distributed import (get_sp_world_size, get_world_rank, get_world_size) from fastvideo.logger import init_logger from fastvideo.models.vision_utils import load_image, load_video @@ -25,26 +24,16 @@ class ValidationDataset(IterableDataset): self.dir = os.path.abspath(self.filename.parent) if not self.filename.exists(): - raise FileNotFoundError( - f"File {self.filename.as_posix()} does not exist") + raise FileNotFoundError(f"File {self.filename.as_posix()} does not exist") if self.filename.suffix == ".csv": - data = datasets.load_dataset("csv", - data_files=self.filename.as_posix(), - split="train") + data = datasets.load_dataset("csv", data_files=self.filename.as_posix(), split="train") elif self.filename.suffix == ".json": - data = datasets.load_dataset("json", - data_files=self.filename.as_posix(), - split="train", - field="data") + data = datasets.load_dataset("json", data_files=self.filename.as_posix(), split="train", field="data") elif self.filename.suffix == ".parquet": - data = datasets.load_dataset("parquet", - data_files=self.filename.as_posix(), - split="train") + data = datasets.load_dataset("parquet", data_files=self.filename.as_posix(), split="train") elif self.filename.suffix == ".arrow": - data = datasets.load_dataset("arrow", - data_files=self.filename.as_posix(), - split="train") + data = datasets.load_dataset("arrow", data_files=self.filename.as_posix(), split="train") else: _SUPPORTED_FILE_FORMATS = [".csv", ".json", ".parquet", ".arrow"] raise ValueError( @@ -69,8 +58,7 @@ class ValidationDataset(IterableDataset): # Duplicate samples cyclically to reach the target additional_samples = [] for i in range(samples_to_add): - additional_samples.append( - self.all_samples[i % self.original_total_samples]) + additional_samples.append(self.all_samples[i % self.original_total_samples]) self.all_samples.extend(additional_samples) @@ -123,8 +111,7 @@ class ValidationDataset(IterableDataset): image_path = sample["image_path"] image_path = os.path.join(self.dir, image_path) sample["image_path"] = image_path - if not pathlib.Path(image_path).is_file( - ) and not image_path.startswith("http"): + if not pathlib.Path(image_path).is_file() and not image_path.startswith("http"): logger.warning("Image file %s does not exist.", image_path) else: sample["image"] = load_image(image_path) @@ -133,8 +120,7 @@ class ValidationDataset(IterableDataset): video_path = sample["video_path"] video_path = os.path.join(self.dir, video_path) sample["video_path"] = video_path - if not pathlib.Path(video_path).is_file( - ) and not video_path.startswith("http"): + if not pathlib.Path(video_path).is_file() and not video_path.startswith("http"): logger.warning("Video file %s does not exist.", video_path) else: sample["video"] = load_video(video_path) @@ -143,8 +129,7 @@ class ValidationDataset(IterableDataset): ref_video_path = sample["ref_video"] ref_video_path = os.path.join(self.dir, ref_video_path) if not pathlib.Path(ref_video_path).is_file(): - logger.warning("Reference video file %s does not exist.", - ref_video_path) + logger.warning("Reference video file %s does not exist.", ref_video_path) sample.pop("ref_video", None) else: sample["ref_video"] = ref_video_path @@ -153,10 +138,8 @@ class ValidationDataset(IterableDataset): control_image_path = sample["control_image_path"] control_image_path = os.path.join(self.dir, control_image_path) sample["control_image_path"] = control_image_path - if not pathlib.Path(control_image_path).is_file( - ) and not control_image_path.startswith("http"): - logger.warning("Control Image file %s does not exist.", - control_image_path) + if not pathlib.Path(control_image_path).is_file() and not control_image_path.startswith("http"): + logger.warning("Control Image file %s does not exist.", control_image_path) else: sample["control_image"] = load_image(control_image_path) @@ -164,10 +147,8 @@ class ValidationDataset(IterableDataset): control_video_path = sample["control_video_path"] control_video_path = os.path.join(self.dir, control_video_path) sample["control_video_path"] = control_video_path - if not pathlib.Path(control_video_path).is_file( - ) and not control_video_path.startswith("http"): - logger.warning("Control Video file %s does not exist.", - control_video_path) + if not pathlib.Path(control_video_path).is_file() and not control_video_path.startswith("http"): + logger.warning("Control Video file %s does not exist.", control_video_path) else: sample["control_video"] = load_video(control_video_path) @@ -176,30 +157,25 @@ class ValidationDataset(IterableDataset): action_path = os.path.join(self.dir, action_path) sample["action_path"] = action_path if not pathlib.Path(action_path).is_file(): - logger.warning("Action file %s does not exist.", - action_path) + logger.warning("Action file %s does not exist.", action_path) else: action_data = np.load(action_path, allow_pickle=True) num_frames = sample["num_frames"] - if isinstance(action_data, - np.ndarray) and action_data.dtype == object: + if isinstance(action_data, np.ndarray) and action_data.dtype == object: action_data = action_data.item() if isinstance(action_data, dict): kb_length = len(action_data["keyboard"]) ms_length = len(action_data["mouse"]) - assert kb_length >= num_frames, ( - f"keyboard length {kb_length} is smaller than " - f"num_frames {num_frames} for {action_path}.") - assert ms_length >= num_frames, ( - f"mouse length {ms_length} is smaller than " - f"num_frames {num_frames} for {action_path}.") + assert kb_length >= num_frames, (f"keyboard length {kb_length} is smaller than " + f"num_frames {num_frames} for {action_path}.") + assert ms_length >= num_frames, (f"mouse length {ms_length} is smaller than " + f"num_frames {num_frames} for {action_path}.") sample["keyboard_cond"] = action_data["keyboard"][:num_frames] sample["mouse_cond"] = action_data["mouse"][:num_frames] else: kb_length = len(action_data) - assert kb_length >= num_frames, ( - f"keyboard length {kb_length} is smaller than " - f"num_frames {num_frames} for {action_path}.") + assert kb_length >= num_frames, (f"keyboard length {kb_length} is smaller than " + f"num_frames {num_frames} for {action_path}.") sample["keyboard_cond"] = action_data[:num_frames] sample = {k: v for k, v in sample.items() if v is not None} diff --git a/fastvideo/models/audio/ltx2_audio_processing.py b/fastvideo/models/audio/ltx2_audio_processing.py index 96dd935b7..9966eb9f8 100644 --- a/fastvideo/models/audio/ltx2_audio_processing.py +++ b/fastvideo/models/audio/ltx2_audio_processing.py @@ -44,8 +44,7 @@ class AudioProcessor(nn.Module): ) -> torch.Tensor: if source_rate == target_rate: return waveform - resampled = torchaudio.functional.resample( - waveform, source_rate, target_rate) + resampled = torchaudio.functional.resample(waveform, source_rate, target_rate) return resampled.to(device=waveform.device, dtype=waveform.dtype) def waveform_to_mel( @@ -53,8 +52,7 @@ class AudioProcessor(nn.Module): waveform: torch.Tensor, waveform_sample_rate: int, ) -> torch.Tensor: - waveform = self.resample_waveform( - waveform, waveform_sample_rate, self.sample_rate) + waveform = self.resample_waveform(waveform, waveform_sample_rate, self.sample_rate) mel = self.mel_transform(waveform) mel = torch.log(torch.clamp(mel, min=1e-5)) mel = mel.to(device=waveform.device, dtype=waveform.dtype) diff --git a/fastvideo/models/audio/ltx2_audio_vae.py b/fastvideo/models/audio/ltx2_audio_vae.py index 28772210c..56df5374c 100644 --- a/fastvideo/models/audio/ltx2_audio_vae.py +++ b/fastvideo/models/audio/ltx2_audio_vae.py @@ -14,7 +14,6 @@ import torch import torch.nn as nn import torch.nn.functional as F - # ============================================================================= # Enums # ============================================================================= @@ -70,11 +69,9 @@ class AudioLatentShape(NamedTuple): # Constants # ============================================================================= - LATENT_DOWNSAMPLE_FACTOR = 4 LRELU_SLOPE = 0.1 - # ============================================================================= # Normalization Layers # ============================================================================= @@ -91,7 +88,7 @@ class PixelNorm(nn.Module): self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: - mean_sq = torch.mean(x ** 2, dim=self.dim, keepdim=True) + mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True) rms = torch.sqrt(mean_sq + self.eps) return x / rms @@ -104,9 +101,7 @@ def build_normalization_layer( ) -> nn.Module: """Create a normalization layer based on the normalization type.""" if normtype == NormType.GROUP: - return nn.GroupNorm( - num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True - ) + return nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True) if normtype == NormType.PIXEL: return PixelNorm(dim=1, eps=1e-6) raise ValueError(f"Invalid normalization type: {normtype}") @@ -128,14 +123,10 @@ class PerChannelStatistics(nn.Module): self.register_buffer("mean-of-means", torch.empty(latent_channels)) def un_normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x * self.get_buffer("std-of-means").to(x)) + self.get_buffer( - "mean-of-means" - ).to(x) + return (x * self.get_buffer("std-of-means").to(x)) + self.get_buffer("mean-of-means").to(x) def normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x - self.get_buffer("mean-of-means").to(x)) / self.get_buffer( - "std-of-means" - ).to(x) + return (x - self.get_buffer("mean-of-means").to(x)) / self.get_buffer("std-of-means").to(x) # ============================================================================= @@ -164,9 +155,7 @@ class AudioPatchifier: """Flatten audio latent tensor along time: (B, C, T, F) -> (B, T, C*F).""" return einops.rearrange(audio_latents, "b c t f -> b t (c f)") - def unpatchify( - self, audio_latents: torch.Tensor, output_shape: AudioLatentShape - ) -> torch.Tensor: + def unpatchify(self, audio_latents: torch.Tensor, output_shape: AudioLatentShape) -> torch.Tensor: """Restore (B, C, T, F) from flattened patches: (B, T, C*F) -> (B, C, T, F).""" return einops.rearrange( audio_latents, @@ -250,9 +239,7 @@ def make_conv2d( ) -> nn.Module: """Create a 2D convolution layer that can be either causal or non-causal.""" if causality_axis is not None: - return CausalConv2d( - in_channels, out_channels, kernel_size, stride, dilation, groups, bias, causality_axis - ) + return CausalConv2d(in_channels, out_channels, kernel_size, stride, dilation, groups, bias, causality_axis) else: if padding is None: padding = kernel_size // 2 if isinstance(kernel_size, int) else tuple(k // 2 for k in kernel_size) @@ -303,7 +290,7 @@ class AttnBlock(nn.Module): q = q.permute(0, 2, 1).contiguous() # b, hw, c k = k.reshape(b, c, h * w).contiguous() # b, c, hw w_ = torch.bmm(q, k).contiguous() # b, hw, hw - w_ = w_ * (int(c) ** (-0.5)) + w_ = w_ * (int(c)**(-0.5)) w_ = F.softmax(w_, dim=2) # Attend to values @@ -365,25 +352,25 @@ class ResnetBlock(nn.Module): self.norm1 = build_normalization_layer(in_channels, normtype=norm_type) self.non_linearity = nn.SiLU() - self.conv1 = make_conv2d( - in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis - ) + self.conv1 = make_conv2d(in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis) if temb_channels > 0: self.temb_proj = nn.Linear(temb_channels, out_channels) self.norm2 = build_normalization_layer(out_channels, normtype=norm_type) self.dropout = nn.Dropout(dropout) - self.conv2 = make_conv2d( - out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis - ) + self.conv2 = make_conv2d(out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis) if self.in_channels != self.out_channels: if self.use_conv_shortcut: - self.conv_shortcut = make_conv2d( - in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis - ) + self.conv_shortcut = make_conv2d(in_channels, + out_channels, + kernel_size=3, + stride=1, + causality_axis=causality_axis) else: - self.nin_shortcut = make_conv2d( - in_channels, out_channels, kernel_size=1, stride=1, causality_axis=causality_axis - ) + self.nin_shortcut = make_conv2d(in_channels, + out_channels, + kernel_size=1, + stride=1, + causality_axis=causality_axis) def forward( self, @@ -418,26 +405,22 @@ class ResBlock1(nn.Module): """1D ResBlock for vocoder with dilated convolutions.""" def __init__( - self, - channels: int, - kernel_size: int = 3, - dilation: Tuple[int, int, int] = (1, 3, 5), + self, + channels: int, + kernel_size: int = 3, + dilation: Tuple[int, int, int] = (1, 3, 5), ): super().__init__() - self.convs1 = nn.ModuleList( - [ - nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding="same"), - nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding="same"), - nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], padding="same"), - ] - ) - self.convs2 = nn.ModuleList( - [ - nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"), - nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"), - nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"), - ] - ) + self.convs1 = nn.ModuleList([ + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], padding="same"), + ]) + self.convs2 = nn.ModuleList([ + nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"), + ]) def forward(self, x: torch.Tensor) -> torch.Tensor: for conv1, conv2 in zip(self.convs1, self.convs2, strict=True): @@ -453,18 +436,16 @@ class ResBlock2(nn.Module): """1D ResBlock for vocoder (simpler version).""" def __init__( - self, - channels: int, - kernel_size: int = 3, - dilation: Tuple[int, int] = (1, 3), + self, + channels: int, + kernel_size: int = 3, + dilation: Tuple[int, int] = (1, 3), ): super().__init__() - self.convs = nn.ModuleList( - [ - nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding="same"), - nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding="same"), - ] - ) + self.convs = nn.ModuleList([ + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding="same"), + ]) def forward(self, x: torch.Tensor) -> torch.Tensor: for conv in self.convs: @@ -491,9 +472,7 @@ def _sinc(x: torch.Tensor) -> torch.Tensor: ) -def kaiser_sinc_filter1d( - cutoff: float, half_width: float, kernel_size: int -) -> torch.Tensor: +def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: even = kernel_size % 2 == 0 half_size = kernel_size // 2 delta_f = 4 * half_width @@ -501,17 +480,11 @@ def kaiser_sinc_filter1d( if amplitude > 50.0: beta = 0.1102 * (amplitude - 8.7) elif amplitude >= 21.0: - beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * ( - amplitude - 21.0 - ) + beta = 0.5842 * (amplitude - 21)**0.4 + 0.07886 * (amplitude - 21.0) else: beta = 0.0 window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) - time = ( - torch.arange(-half_size, half_size) + 0.5 - if even - else torch.arange(kernel_size) - half_size - ) + time = (torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size) if cutoff == 0: filter_ = torch.zeros_like(time) else: @@ -521,6 +494,7 @@ def kaiser_sinc_filter1d( class LowPassFilter1d(nn.Module): + def __init__( self, cutoff: float = 0.5, @@ -556,6 +530,7 @@ class LowPassFilter1d(nn.Module): class UpSample1d(nn.Module): + def __init__( self, ratio: int = 2, @@ -576,26 +551,14 @@ class UpSample1d(nn.Module): self.pad_left = 2 * width * ratio self.pad_right = self.kernel_size - ratio time_axis = (torch.arange(self.kernel_size) / ratio - width) * rolloff - time_clamped = time_axis.clamp( - -lowpass_filter_width, lowpass_filter_width - ) - window = ( - torch.cos(time_clamped * math.pi / lowpass_filter_width / 2) ** 2 - ) - sinc_filter = ( - torch.sinc(time_axis) * window * rolloff / ratio - ).view(1, 1, -1) + time_clamped = time_axis.clamp(-lowpass_filter_width, lowpass_filter_width) + window = (torch.cos(time_clamped * math.pi / lowpass_filter_width / 2)**2) + sinc_filter = (torch.sinc(time_axis) * window * rolloff / ratio).view(1, 1, -1) else: - self.kernel_size = ( - int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size - ) + self.kernel_size = (int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size) self.pad = self.kernel_size // ratio - 1 - self.pad_left = self.pad * self.stride + ( - self.kernel_size - self.stride - ) // 2 - self.pad_right = self.pad * self.stride + ( - self.kernel_size - self.stride + 1 - ) // 2 + self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 sinc_filter = kaiser_sinc_filter1d( cutoff=0.5 / ratio, half_width=0.6 / ratio, @@ -607,22 +570,17 @@ class UpSample1d(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: _, n_channels, _ = x.shape x = F.pad(x, (self.pad, self.pad), mode="replicate") - filt = self.filter.to(dtype=x.dtype, device=x.device).expand( - n_channels, -1, -1 - ) - x = self.ratio * F.conv_transpose1d( - x, filt, stride=self.stride, groups=n_channels - ) - return x[..., self.pad_left : -self.pad_right] + filt = self.filter.to(dtype=x.dtype, device=x.device).expand(n_channels, -1, -1) + x = self.ratio * F.conv_transpose1d(x, filt, stride=self.stride, groups=n_channels) + return x[..., self.pad_left:-self.pad_right] class DownSample1d(nn.Module): + def __init__(self, ratio: int = 2, kernel_size: int | None = None) -> None: super().__init__() self.ratio = ratio - self.kernel_size = ( - int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size - ) + self.kernel_size = (int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size) self.lowpass = LowPassFilter1d( cutoff=0.5 / ratio, half_width=0.6 / ratio, @@ -635,6 +593,7 @@ class DownSample1d(nn.Module): class Activation1d(nn.Module): + def __init__( self, activation: nn.Module, @@ -655,6 +614,7 @@ class Activation1d(nn.Module): class Snake(nn.Module): + def __init__( self, in_features: int, @@ -664,11 +624,7 @@ class Snake(nn.Module): ) -> None: super().__init__() self.alpha_logscale = alpha_logscale - self.alpha = nn.Parameter( - torch.zeros(in_features) - if alpha_logscale - else torch.ones(in_features) * alpha - ) + self.alpha = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha) self.alpha.requires_grad = alpha_trainable self.eps = 1e-9 @@ -680,6 +636,7 @@ class Snake(nn.Module): class SnakeBeta(nn.Module): + def __init__( self, in_features: int, @@ -689,17 +646,9 @@ class SnakeBeta(nn.Module): ) -> None: super().__init__() self.alpha_logscale = alpha_logscale - self.alpha = nn.Parameter( - torch.zeros(in_features) - if alpha_logscale - else torch.ones(in_features) * alpha - ) + self.alpha = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha) self.alpha.requires_grad = alpha_trainable - self.beta = nn.Parameter( - torch.zeros(in_features) - if alpha_logscale - else torch.ones(in_features) * alpha - ) + self.beta = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha) self.beta.requires_grad = alpha_trainable self.eps = 1e-9 @@ -713,82 +662,73 @@ class SnakeBeta(nn.Module): class AMPBlock1(nn.Module): + def __init__( - self, - channels: int, - kernel_size: int = 3, - dilation: tuple[int, int, int] = (1, 3, 5), - activation: str = "snake", + self, + channels: int, + kernel_size: int = 3, + dilation: tuple[int, int, int] = (1, 3, 5), + activation: str = "snake", ) -> None: super().__init__() act_cls = SnakeBeta if activation == "snakebeta" else Snake - self.convs1 = nn.ModuleList( - [ - nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[0], - padding=get_padding(kernel_size, dilation[0]), - ), - nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[1], - padding=get_padding(kernel_size, dilation[1]), - ), - nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[2], - padding=get_padding(kernel_size, dilation[2]), - ), - ] - ) - self.convs2 = nn.ModuleList( - [ - nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=1, - padding=get_padding(kernel_size, 1), - ), - nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=1, - padding=get_padding(kernel_size, 1), - ), - nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=1, - padding=get_padding(kernel_size, 1), - ), - ] - ) - self.acts1 = nn.ModuleList( - [Activation1d(act_cls(channels)) for _ in range(len(self.convs1))] - ) - self.acts2 = nn.ModuleList( - [Activation1d(act_cls(channels)) for _ in range(len(self.convs2))] - ) + self.convs1 = nn.ModuleList([ + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[2], + padding=get_padding(kernel_size, dilation[2]), + ), + ]) + self.convs2 = nn.ModuleList([ + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ), + ]) + self.acts1 = nn.ModuleList([Activation1d(act_cls(channels)) for _ in range(len(self.convs1))]) + self.acts2 = nn.ModuleList([Activation1d(act_cls(channels)) for _ in range(len(self.convs2))]) def forward(self, x: torch.Tensor) -> torch.Tensor: - for c1, c2, a1, a2 in zip( - self.convs1, self.convs2, self.acts1, self.acts2, strict=True - ): + for c1, c2, a1, a2 in zip(self.convs1, self.convs2, self.acts1, self.acts2, strict=True): xt = a1(x) xt = c1(xt) xt = a2(xt) @@ -879,8 +819,7 @@ def build_downsampling_path( dropout=dropout, norm_type=norm_type, causality_axis=causality_axis, - ) - ) + )) block_in = block_out if curr_res in attn_resolutions: attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type)) @@ -914,9 +853,7 @@ class Upsample(nn.Module): self.with_conv = with_conv self.causality_axis = causality_axis if self.with_conv: - self.conv = make_conv2d( - in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis - ) + self.conv = make_conv2d(in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis) def forward(self, x: torch.Tensor) -> torch.Tensor: x = F.interpolate(x, scale_factor=2.0, mode="nearest") @@ -956,7 +893,7 @@ def build_upsampling_path( """Build the upsampling path with residual blocks, attention, and upsampling layers.""" up_modules = nn.ModuleList() block_in = initial_block_channels - curr_res = resolution // (2 ** (num_resolutions - 1)) + curr_res = resolution // (2**(num_resolutions - 1)) for level in reversed(range(num_resolutions)): stage = nn.Module() @@ -973,8 +910,7 @@ def build_upsampling_path( dropout=dropout, norm_type=norm_type, causality_axis=causality_axis, - ) - ) + )) block_in = block_out if curr_res in attn_resolutions: stage.attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type)) @@ -1012,11 +948,7 @@ def build_mid_block( norm_type=norm_type, causality_axis=causality_axis, ) - mid.attn_1 = ( - make_attn(channels, attn_type=attn_type, norm_type=norm_type) - if add_attention - else nn.Identity() - ) + mid.attn_1 = (make_attn(channels, attn_type=attn_type, norm_type=norm_type) if add_attention else nn.Identity()) mid.block_2 = ResnetBlock( in_channels=channels, out_channels=channels, @@ -1258,12 +1190,14 @@ class AudioDecoder(nn.Module): self.attn_type = attn_type base_block_channels = ch * self.channel_multipliers[-1] - base_resolution = resolution // (2 ** (self.num_resolutions - 1)) + base_resolution = resolution // (2**(self.num_resolutions - 1)) self.z_shape = (1, z_channels, base_resolution, base_resolution) - self.conv_in = make_conv2d( - z_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis - ) + self.conv_in = make_conv2d(z_channels, + base_block_channels, + kernel_size=3, + stride=1, + causality_axis=self.causality_axis) self.non_linearity = nn.SiLU() self.mid = build_mid_block( @@ -1293,9 +1227,11 @@ class AudioDecoder(nn.Module): ) self.norm_out = build_normalization_layer(final_block_channels, normtype=self.norm_type) - self.conv_out = make_conv2d( - final_block_channels, out_ch, kernel_size=3, stride=1, causality_axis=self.causality_axis - ) + self.conv_out = make_conv2d(final_block_channels, + out_ch, + kernel_size=3, + stride=1, + causality_axis=self.causality_axis) def forward(self, sample: torch.Tensor) -> torch.Tensor: """ @@ -1314,9 +1250,7 @@ class AudioDecoder(nn.Module): return self._adjust_output_shape(h, target_shape) - def _denormalize_latents( - self, sample: torch.Tensor - ) -> Tuple[torch.Tensor, AudioLatentShape]: + def _denormalize_latents(self, sample: torch.Tensor) -> Tuple[torch.Tensor, AudioLatentShape]: latent_shape = AudioLatentShape( batch=sample.shape[0], channels=sample.shape[1], @@ -1353,9 +1287,8 @@ class AudioDecoder(nn.Module): target_freq = target_shape.mel_bins # Crop first - decoded_output = decoded_output[ - :, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq) - ] + decoded_output = decoded_output[:, :target_channels, :min(current_time, target_time + ), :min(current_freq, target_freq)] # Calculate padding needed time_padding_needed = target_time - decoded_output.shape[2] @@ -1456,21 +1389,18 @@ class Vocoder(nn.Module): self.ups.append( nn.ConvTranspose1d( upsample_initial_channel // (2**i), - upsample_initial_channel // (2 ** (i + 1)), + upsample_initial_channel // (2**(i + 1)), kernel_size, stride, padding=(kernel_size - stride) // 2, - ) - ) + )) self.resblocks = nn.ModuleList() for i, _ in enumerate(self.ups): - ch = upsample_initial_channel // (2 ** (i + 1)) + ch = upsample_initial_channel // (2**(i + 1)) for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True): if self.is_amp: - self.resblocks.append( - resblock_class(ch, kernel_size, dilations, activation=activation) - ) + self.resblocks.append(resblock_class(ch, kernel_size, dilations, activation=activation)) else: self.resblocks.append(resblock_class(ch, kernel_size, dilations)) @@ -1534,8 +1464,7 @@ class Vocoder(nn.Module): # ============================================================================= -def _build_stft_basis(filter_length: int, - win_length: int) -> torch.Tensor: +def _build_stft_basis(filter_length: int, win_length: int) -> torch.Tensor: """Hann-windowed FFT basis used by ``_STFTFn`` as a Conv1d kernel. Returns a ``(2*(filter_length//2 + 1), 1, filter_length)`` float32 tensor @@ -1552,9 +1481,7 @@ def _build_stft_basis(filter_length: int, # FFT matrix: row k is e^{-2*pi*j*k*n / filter_length} for n in [0, filter_length) fft_matrix = torch.fft.fft(torch.eye(filter_length, dtype=torch.float32)) fft_basis = fft_matrix[:n_freqs] # (n_freqs, filter_length) complex - window = torch.hann_window(win_length, - periodic=True, - dtype=torch.float32) + window = torch.hann_window(win_length, periodic=True, dtype=torch.float32) if win_length < filter_length: pad_left = (filter_length - win_length) // 2 pad_right = filter_length - win_length - pad_left @@ -1584,10 +1511,7 @@ def _build_mel_basis(sampling_rate: int, if fmax is None: fmax = sampling_rate / 2.0 n_freqs = n_fft // 2 + 1 - fft_freqs = torch.linspace(0.0, - sampling_rate / 2.0, - n_freqs, - dtype=torch.float32) + fft_freqs = torch.linspace(0.0, sampling_rate / 2.0, n_freqs, dtype=torch.float32) def _hz_to_mel(hz: float) -> float: return 2595.0 * math.log10(1.0 + hz / 700.0) @@ -1597,11 +1521,7 @@ def _build_mel_basis(sampling_rate: int, mel_lo = _hz_to_mel(fmin) mel_hi = _hz_to_mel(fmax) - mel_edges_hz = _mel_to_hz( - torch.linspace(mel_lo, - mel_hi, - n_mel_channels + 2, - dtype=torch.float32)) + mel_edges_hz = _mel_to_hz(torch.linspace(mel_lo, mel_hi, n_mel_channels + 2, dtype=torch.float32)) weights = torch.zeros(n_mel_channels, n_freqs, dtype=torch.float32) for i in range(n_mel_channels): @@ -1611,12 +1531,12 @@ def _build_mel_basis(sampling_rate: int, lower = (fft_freqs - f_lo) / (f_ctr - f_lo) upper = (f_hi - fft_freqs) / (f_hi - f_ctr) weights[i] = torch.clamp(torch.minimum(lower, upper), min=0.0) - enorm = 2.0 / (mel_edges_hz[2:n_mel_channels + 2] - - mel_edges_hz[:n_mel_channels]) + enorm = 2.0 / (mel_edges_hz[2:n_mel_channels + 2] - mel_edges_hz[:n_mel_channels]) return weights * enorm.unsqueeze(1) class _STFTFn(nn.Module): + def __init__(self, filter_length: int, hop_length: int, win_length: int) -> None: super().__init__() self.hop_length = hop_length @@ -1643,6 +1563,7 @@ class _STFTFn(nn.Module): class MelSTFT(nn.Module): + def __init__( self, filter_length: int, @@ -1660,9 +1581,7 @@ class MelSTFT(nn.Module): _build_mel_basis(sampling_rate, filter_length, n_mel_channels), ) - def mel_spectrogram( - self, y: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + def mel_spectrogram(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: magnitude, phase = self.stft_fn(y) energy = torch.norm(magnitude, dim=1) mel = torch.matmul(self.mel_basis.to(magnitude.dtype), magnitude) @@ -1671,6 +1590,7 @@ class MelSTFT(nn.Module): class VocoderWithBWE(nn.Module): + def __init__( self, vocoder: Vocoder, @@ -1711,11 +1631,7 @@ class VocoderWithBWE(nn.Module): def forward(self, mel_spec: torch.Tensor) -> torch.Tensor: x = self.vocoder(mel_spec) _, _, length_low_rate = x.shape - output_length = ( - length_low_rate - * self.output_sampling_rate - // self.input_sampling_rate - ) + output_length = (length_low_rate * self.output_sampling_rate // self.input_sampling_rate) remainder = length_low_rate % self.hop_length if remainder != 0: @@ -1725,9 +1641,7 @@ class VocoderWithBWE(nn.Module): mel_for_bwe = mel.transpose(2, 3) residual = self.bwe_generator(mel_for_bwe) skip = self.resampler(x) - assert residual.shape == skip.shape, ( - f"residual {residual.shape} != skip {skip.shape}" - ) + assert residual.shape == skip.shape, (f"residual {residual.shape} != skip {skip.shape}") return torch.clamp(residual + skip, -1, 1)[..., :output_length] @@ -1827,16 +1741,12 @@ class VocoderConfigurator: nested_vocoder_cfg = vocoder_cfg.get("vocoder", {}) bwe_cfg = vocoder_cfg.get("bwe", {}) base_vocoder = Vocoder( - resblock_kernel_sizes=nested_vocoder_cfg.get( - "resblock_kernel_sizes", [3, 7, 11]), - upsample_rates=nested_vocoder_cfg.get( - "upsample_rates", [6, 5, 2, 2, 2]), - upsample_kernel_sizes=nested_vocoder_cfg.get( - "upsample_kernel_sizes", [16, 15, 8, 4, 4]), - resblock_dilation_sizes=nested_vocoder_cfg.get( - "resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]), - upsample_initial_channel=nested_vocoder_cfg.get( - "upsample_initial_channel", 1024), + resblock_kernel_sizes=nested_vocoder_cfg.get("resblock_kernel_sizes", [3, 7, 11]), + upsample_rates=nested_vocoder_cfg.get("upsample_rates", [6, 5, 2, 2, 2]), + upsample_kernel_sizes=nested_vocoder_cfg.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]), + resblock_dilation_sizes=nested_vocoder_cfg.get("resblock_dilation_sizes", + [[1, 3, 5], [1, 3, 5], [1, 3, 5]]), + upsample_initial_channel=nested_vocoder_cfg.get("upsample_initial_channel", 1024), stereo=nested_vocoder_cfg.get("stereo", True), resblock=nested_vocoder_cfg.get("resblock", "AMP1"), output_sample_rate=bwe_cfg.get( @@ -1844,31 +1754,22 @@ class VocoderConfigurator: nested_vocoder_cfg.get("output_sampling_rate", 24000), ), activation=nested_vocoder_cfg.get("activation", "snakebeta"), - use_tanh_at_final=nested_vocoder_cfg.get( - "use_tanh_at_final", True), - apply_final_activation=nested_vocoder_cfg.get( - "apply_final_activation", True), - use_bias_at_final=nested_vocoder_cfg.get( - "use_bias_at_final", True), + use_tanh_at_final=nested_vocoder_cfg.get("use_tanh_at_final", True), + apply_final_activation=nested_vocoder_cfg.get("apply_final_activation", True), + use_bias_at_final=nested_vocoder_cfg.get("use_bias_at_final", True), ) bwe_generator = Vocoder( - resblock_kernel_sizes=bwe_cfg.get( - "resblock_kernel_sizes", [3, 7, 11]), - upsample_rates=bwe_cfg.get( - "upsample_rates", [6, 5, 2, 2, 2]), - upsample_kernel_sizes=bwe_cfg.get( - "upsample_kernel_sizes", [16, 15, 8, 4, 4]), - resblock_dilation_sizes=bwe_cfg.get( - "resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]), - upsample_initial_channel=bwe_cfg.get( - "upsample_initial_channel", 1024), + resblock_kernel_sizes=bwe_cfg.get("resblock_kernel_sizes", [3, 7, 11]), + upsample_rates=bwe_cfg.get("upsample_rates", [6, 5, 2, 2, 2]), + upsample_kernel_sizes=bwe_cfg.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]), + resblock_dilation_sizes=bwe_cfg.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]), + upsample_initial_channel=bwe_cfg.get("upsample_initial_channel", 1024), stereo=bwe_cfg.get("stereo", True), resblock=bwe_cfg.get("resblock", "AMP1"), output_sample_rate=bwe_cfg.get("output_sampling_rate", 24000), activation=bwe_cfg.get("activation", "snakebeta"), use_tanh_at_final=bwe_cfg.get("use_tanh_at_final", True), - apply_final_activation=bwe_cfg.get( - "apply_final_activation", False), + apply_final_activation=bwe_cfg.get("apply_final_activation", False), use_bias_at_final=bwe_cfg.get("use_bias_at_final", True), ) mel_stft = MelSTFT( @@ -1941,9 +1842,7 @@ class LTX2Vocoder(nn.Module): return self.model(spectrogram) -def decode_audio( - latent: torch.Tensor, audio_decoder: AudioDecoder, vocoder: Vocoder -) -> torch.Tensor: +def decode_audio(latent: torch.Tensor, audio_decoder: AudioDecoder, vocoder: Vocoder) -> torch.Tensor: """ Decode an audio latent representation using the provided audio decoder and vocoder. """ @@ -1951,5 +1850,6 @@ def decode_audio( decoded_audio = vocoder(decoded_audio).squeeze(0).float() return decoded_audio + # Entry point for model registry EntryClass = [LTX2AudioEncoder, LTX2AudioDecoder, LTX2Vocoder] diff --git a/fastvideo/models/camera/trajectory.py b/fastvideo/models/camera/trajectory.py index bc3b7abf1..796d363b8 100644 --- a/fastvideo/models/camera/trajectory.py +++ b/fastvideo/models/camera/trajectory.py @@ -14,7 +14,6 @@ import numpy as np import torch from packaging import version as pver - # Action name -> motion type mapping ACTION_DICT = { "w": "forward", @@ -175,9 +174,7 @@ def _action_to_pose_list(action_id: str, value: float = 0.2, duration: int = 33) intrinsic = [0.50505, 0.8979, 0.5, 0.5] motion_type = ACTION_DICT.get(action_id, action_id) - positions, rotations, current_pose = _generate_motion_segment( - current_pose, motion_type, value, duration - ) + positions, rotations, current_pose = _generate_motion_segment(current_pose, motion_type, value, duration) all_positions.extend(positions) all_rotations.extend(rotations) @@ -347,9 +344,7 @@ def create_camera_trajectory( # Convert to w2c matrices w2cs = [np.asarray([float(p) for p in pose[7:]]).reshape(3, 4) for pose in poses_parsed] - transform_matrix = np.asarray( - [[1, 0, 0, 0], [0, 0, 1, 0], [0, -1, 0, 0], [0, 0, 0, 1]] - ).reshape(4, 4) + transform_matrix = np.asarray([[1, 0, 0, 0], [0, 0, 1, 0], [0, -1, 0, 0], [0, 0, 0, 1]]).reshape(4, 4) last_row = np.zeros((1, 4)) last_row[0, -1] = 1.0 w2cs = [np.concatenate((w2c, last_row), axis=0) for w2c in w2cs] @@ -365,15 +360,12 @@ def create_camera_trajectory( monst3r_h = cam_params[0].cy * 2 ratio_w, ratio_h = width / monst3r_w, height / monst3r_h intrinsics = np.asarray( - [ - [ - cam_param.fx * ratio_w, - cam_param.fy * ratio_h, - cam_param.cx * ratio_w, - cam_param.cy * ratio_h, - ] - for cam_param in cam_params - ], + [[ + cam_param.fx * ratio_w, + cam_param.fy * ratio_h, + cam_param.cx * ratio_w, + cam_param.cy * ratio_h, + ] for cam_param in cam_params], dtype=np.float32, ) intrinsics = torch.as_tensor(intrinsics)[None] # [1, n_frame, 4] diff --git a/fastvideo/models/dits/base.py b/fastvideo/models/dits/base.py index 5a9b1d276..61929f6d2 100644 --- a/fastvideo/models/dits/base.py +++ b/fastvideo/models/dits/base.py @@ -19,30 +19,21 @@ class BaseDiT(nn.Module, ABC): num_attention_heads: int num_channels_latents: int # always supports torch_sdpa - _supported_attention_backends: tuple[ - AttentionBackendEnum, ...] = DiTConfig()._supported_attention_backends + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = DiTConfig()._supported_attention_backends def __init_subclass__(cls) -> None: - required_class_attrs = [ - "_fsdp_shard_conditions", "param_names_mapping", - "_compile_conditions" - ] + required_class_attrs = ["_fsdp_shard_conditions", "param_names_mapping", "_compile_conditions"] super().__init_subclass__() for attr in required_class_attrs: if not hasattr(cls, attr): - raise AttributeError( - f"Subclasses of BaseDiT must define '{attr}' class variable" - ) + raise AttributeError(f"Subclasses of BaseDiT must define '{attr}' class variable") - def __init__(self, config: DiTConfig, hf_config: dict[str, Any], - **kwargs) -> None: + def __init__(self, config: DiTConfig, hf_config: dict[str, Any], **kwargs) -> None: super().__init__() self.config = config self.hf_config = hf_config if not self.supported_attention_backends: - raise ValueError( - f"Subclass {self.__class__.__name__} must define _supported_attention_backends" - ) + raise ValueError(f"Subclass {self.__class__.__name__} must define _supported_attention_backends") @abstractmethod def forward(self, @@ -56,14 +47,10 @@ class BaseDiT(nn.Module, ABC): pass def __post_init__(self) -> None: - required_attrs = [ - "hidden_size", "num_attention_heads", "num_channels_latents" - ] + required_attrs = ["hidden_size", "num_attention_heads", "num_channels_latents"] for attr in required_attrs: if not hasattr(self, attr): - raise AttributeError( - f"Subclasses of BaseDiT must define '{attr}' instance variable" - ) + raise AttributeError(f"Subclasses of BaseDiT must define '{attr}' instance variable") @property def supported_attention_backends(self) -> tuple[AttentionBackendEnum, ...]: diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 015615668..2875edec1 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -12,23 +12,19 @@ from torch.nn.attention.flex_attention import BlockMask # wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention # see https://github.com/pytorch/pytorch/issues/133254 # change to default for other models -flex_attention = torch.compile( - flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") +flex_attention = torch.compile(flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") import torch.distributed as dist import fastvideo.envs as envs -from fastvideo.attention import (DistributedAttention, - LocalAttention) +from fastvideo.attention import (DistributedAttention, LocalAttention) from fastvideo.configs.models.dits import WanVideoConfig from fastvideo.distributed.parallel_state import get_sp_world_size from fastvideo.forward_context import get_forward_context -from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, - RMSNorm, ScaleResidual, +from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual, ScaleResidualLayerNormScaleShift) from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.mlp import MLP -from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, - get_rotary_pos_embed) +from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, get_rotary_pos_embed) from fastvideo.layers.visual_embedding import (PatchEmbed) from fastvideo.logger import init_logger from fastvideo.models.dits._relative_rope import relativistic_window_offsets @@ -65,16 +61,15 @@ class CausalWanSelfAttention(nn.Module): self.rope_cache_policy = rope_cache_policy # Scaled dot product attention - self.attn = LocalAttention( - num_heads=num_heads, - head_size=self.head_dim, - dropout_rate=0, - softmax_scale=None, - causal=False, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA)) + self.attn = LocalAttention(num_heads=num_heads, + head_size=self.head_dim, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA)) - def forward(self, + def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -106,31 +101,26 @@ class CausalWanSelfAttention(nn.Module): if kv_cache is None: # Padding for flex attention padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] - padded_roped_query = torch.cat( - [roped_query, - torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], - device=q.device, dtype=v.dtype)], - dim=1 - ) + padded_roped_query = torch.cat([ + roped_query, + torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], device=q.device, dtype=v.dtype) + ], + dim=1) - padded_roped_key = torch.cat( - [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], - device=k.device, dtype=v.dtype)], - dim=1 - ) + padded_roped_key = torch.cat([ + roped_key, + torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], device=k.device, dtype=v.dtype) + ], + dim=1) padded_v = torch.cat( - [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], - device=v.device, dtype=v.dtype)], - dim=1 - ) + [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], device=v.device, dtype=v.dtype)], + dim=1) - x = flex_attention( - query=padded_roped_query.transpose(2, 1), - key=padded_roped_key.transpose(2, 1), - value=padded_v.transpose(2, 1), - block_mask=block_mask - )[:, :, :-padded_length].transpose(2, 1) + x = flex_attention(query=padded_roped_query.transpose(2, 1), + key=padded_roped_key.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask)[:, :, :-padded_length].transpose(2, 1) else: current_end = current_start + q.shape[1] sink_tokens = self.sink_size * frame_seqlen @@ -139,28 +129,22 @@ class CausalWanSelfAttention(nn.Module): else: max_attention_size = self.local_attn_size * frame_seqlen if self.local_attn_size == -1 and current_end > max_attention_size: - raise ValueError( - "Causal Wan local_attn_size=-1 keeps the previous " - f"{GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES}-latent-frame KV " - "window for compatibility. Set local_attn_size for " - f"longer rollouts; got current_end={current_end} tokens " - f"with frame_seqlen={frame_seqlen}.") + raise ValueError("Causal Wan local_attn_size=-1 keeps the previous " + f"{GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES}-latent-frame KV " + "window for compatibility. Set local_attn_size for " + f"longer rollouts; got current_end={current_end} tokens " + f"with frame_seqlen={frame_seqlen}.") # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache kv_cache_size = kv_cache["k"].shape[1] num_new_tokens = q.shape[1] stored_key = k if relativistic else roped_key # raw vs roped in cache - global_end_index = ( - int(kv_cache["global_end_index"].item()) - if isinstance(kv_cache["global_end_index"], torch.Tensor) - else int(kv_cache["global_end_index"]) - ) - local_end_index_prev = ( - int(kv_cache["local_end_index"].item()) - if isinstance(kv_cache["local_end_index"], torch.Tensor) - else int(kv_cache["local_end_index"]) - ) - if self.local_attn_size != -1 and (current_end > global_end_index) and ( - num_new_tokens + local_end_index_prev > kv_cache_size): + global_end_index = (int(kv_cache["global_end_index"].item()) if isinstance( + kv_cache["global_end_index"], torch.Tensor) else int(kv_cache["global_end_index"])) + local_end_index_prev = (int(kv_cache["local_end_index"].item()) if isinstance( + kv_cache["local_end_index"], torch.Tensor) else int(kv_cache["local_end_index"])) + if self.local_attn_size != -1 and (current_end + > global_end_index) and (num_new_tokens + local_end_index_prev + > kv_cache_size): # Calculate the number of new tokens added in this step # Shift existing cache content left to discard oldest tokens # Clone the source slice to avoid overlapping memory error @@ -188,14 +172,12 @@ class CausalWanSelfAttention(nn.Module): key_window = kv_cache["k"][:, max(0, local_end_index - max_attention_size):local_end_index] value_window = kv_cache["v"][:, max(0, local_end_index - max_attention_size):local_end_index] if relativistic: - window_len, query_lo, query_hi = relativistic_window_offsets( - local_end_index, num_new_tokens, max_attention_size) - roped_query = _apply_rotary_emb( - q, cos[query_lo:query_hi], sin[query_lo:query_hi], - is_neox_style=False).type_as(v) - key_window = _apply_rotary_emb( - key_window, cos[:window_len], sin[:window_len], - is_neox_style=False).type_as(v) + window_len, query_lo, query_hi = relativistic_window_offsets(local_end_index, num_new_tokens, + max_attention_size) + roped_query = _apply_rotary_emb(q, cos[query_lo:query_hi], sin[query_lo:query_hi], + is_neox_style=False).type_as(v) + key_window = _apply_rotary_emb(key_window, cos[:window_len], sin[:window_len], + is_neox_style=False).type_as(v) x = self.attn(roped_query, key_window, value_window) if isinstance(kv_cache["global_end_index"], torch.Tensor): kv_cache["global_end_index"].fill_(current_end) @@ -208,6 +190,7 @@ class CausalWanSelfAttention(nn.Module): return x + class CausalWanTransformerBlock(nn.Module): def __init__(self, @@ -232,14 +215,13 @@ class CausalWanTransformerBlock(nn.Module): self.to_v = ReplicatedLinear(dim, dim, bias=True) self.to_out = ReplicatedLinear(dim, dim, bias=True) - self.attn1 = CausalWanSelfAttention( - dim, - num_heads, - local_attn_size=local_attn_size, - sink_size=sink_size, - qk_norm=qk_norm, - eps=eps, - rope_cache_policy=rope_cache_policy) + self.attn1 = CausalWanSelfAttention(dim, + num_heads, + local_attn_size=local_attn_size, + sink_size=sink_size, + qk_norm=qk_norm, + eps=eps, + rope_cache_policy=rope_cache_policy) self.hidden_dim = dim self.num_attention_heads = num_heads self.local_attn_size = local_attn_size @@ -255,25 +237,20 @@ class CausalWanTransformerBlock(nn.Module): print("QK Norm type not supported") raise Exception assert cross_attn_norm is True - self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( - dim, - norm_type="layer", - eps=eps, - elementwise_affine=True, - dtype=torch.float32) + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32) # 2. Cross-attention # Only T2V for now - self.attn2 = WanT2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps) - self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( - dim, - norm_type="layer", - eps=eps, - elementwise_affine=False, - dtype=torch.float32) + self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32) # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") @@ -310,14 +287,13 @@ class CausalWanTransformerBlock(nn.Module): e = self.scale_shift_table + temb # e.shape: [batch_size, temb_seq_len, 6, inner_dim] assert e.shape == (bs, temb_seq_len, 6, self.hidden_dim) - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( - 6, dim=2) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=2) # *_msa.shape: [batch_size, temb_seq_len, 1, inner_dim] # assert shift_msa.dtype == torch.float32 # 1. Self-attention norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(temb_seq_len, tokens_per_temb)) * - (1 + scale_msa) + shift_msa).flatten(1, 2) + (1 + scale_msa) + shift_msa).flatten(1, 2) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -347,18 +323,17 @@ class CausalWanTransformerBlock(nn.Module): attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.tensor([0], device=hidden_states.device) - norm_hidden_states, hidden_states = self.self_attn_residual_norm( - hidden_states, attn_output, gate_msa, null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, + null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) # 2. Cross-attention attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None, crossattn_cache=crossattn_cache) - norm_hidden_states, hidden_states = self.cross_attn_residual_norm( - hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) + norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa, + c_scale_msa) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) @@ -366,17 +341,16 @@ class CausalWanTransformerBlock(nn.Module): return hidden_states + class CausalWanTransformer3DModel(BaseDiT): _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions _compile_conditions = WanVideoConfig()._compile_conditions - _supported_attention_backends = WanVideoConfig( - )._supported_attention_backends + _supported_attention_backends = WanVideoConfig()._supported_attention_backends param_names_mapping = WanVideoConfig().param_names_mapping reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping - def __init__(self, config: WanVideoConfig, hf_config: dict[str, - Any]) -> None: + def __init__(self, config: WanVideoConfig, hf_config: dict[str, Any]) -> None: super().__init__(config=config, hf_config=hf_config) inner_dim = config.num_attention_heads * config.attention_head_dim @@ -408,17 +382,17 @@ class CausalWanTransformer3DModel(BaseDiT): # 3. Transformer blocks self.blocks = nn.ModuleList([ CausalWanTransformerBlock(inner_dim, - config.ffn_dim, - config.num_attention_heads, - config.local_attn_size, - config.sink_size, - config.qk_norm, - config.cross_attn_norm, - config.eps, - config.added_kv_proj_dim, - self._supported_attention_backends, - prefix=f"{config.prefix}.blocks.{i}", - rope_cache_policy=config.arch_config.rope_cache_policy) + config.ffn_dim, + config.num_attention_heads, + config.local_attn_size, + config.sink_size, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + self._supported_attention_backends, + prefix=f"{config.prefix}.blocks.{i}", + rope_cache_policy=config.arch_config.rope_cache_policy) for i in range(config.num_layers) ]) @@ -428,10 +402,8 @@ class CausalWanTransformer3DModel(BaseDiT): eps=config.eps, elementwise_affine=False, dtype=torch.float32) - self.proj_out = nn.Linear( - inner_dim, config.out_channels * math.prod(config.patch_size)) - self.scale_shift_table = nn.Parameter( - torch.randn(1, 2, inner_dim) / inner_dim**0.5) + self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) self.gradient_checkpointing = False @@ -445,10 +417,11 @@ class CausalWanTransformer3DModel(BaseDiT): self.__post_init__() @staticmethod - def _prepare_blockwise_causal_attn_mask( - device: torch.device | str, num_frames: int = 21, - frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1 - ) -> BlockMask: + def _prepare_blockwise_causal_attn_mask(device: torch.device | str, + num_frames: int = 21, + frame_seqlen: int = 1560, + num_frame_per_block=1, + local_attn_size=-1) -> BlockMask: """ we will divide the token sequence into the following format [1 latent frame] [1 latent frame] ... [1 latent frame] @@ -459,16 +432,10 @@ class CausalWanTransformer3DModel(BaseDiT): # we do right padding to get to a multiple of 128 padded_length = math.ceil(total_length / 128) * 128 - total_length - ends = torch.zeros(total_length + padded_length, - device=device, dtype=torch.long) + ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) # Block-wise causal mask will attend to all elements that are before the end of the current chunk - frame_indices = torch.arange( - start=0, - end=total_length, - step=frame_seqlen * num_frame_per_block, - device=device - ) + frame_indices = torch.arange(start=0, end=total_length, step=frame_seqlen * num_frame_per_block, device=device) for tmp in frame_indices: ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ @@ -478,15 +445,20 @@ class CausalWanTransformer3DModel(BaseDiT): if local_attn_size == -1: return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) else: - return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) + return ((kv_idx < ends[q_idx]) & + (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask - block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, - KV_LEN=total_length + padded_length, _compile=False, device=device) + block_mask = create_block_mask(attention_mask, + B=None, + H=None, + Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, + _compile=False, + device=device) if not dist.is_initialized() or dist.get_rank() == 0: - print( - f" cache a block wise causal mask with block size of {num_frame_per_block} frames") + print(f" cache a block wise causal mask with block size of {num_frame_per_block} frames") print(block_mask) # import imageio @@ -502,22 +474,22 @@ class CausalWanTransformer3DModel(BaseDiT): return block_mask @staticmethod - def _prepare_teacher_forcing_mask( - device: torch.device | str, num_frames: int = 21, - frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1 - ) -> BlockMask: + def _prepare_teacher_forcing_mask(device: torch.device | str, + num_frames: int = 21, + frame_seqlen: int = 1560, + num_frame_per_block=1, + local_attn_size=-1) -> BlockMask: """Attention mask for the teacher-forcing ``[clean | noisy]`` sequence. A noisy token attends to its own block plus the clean context of all strictly previous blocks; clean tokens are block-wise causal. """ if local_attn_size != -1: - raise NotImplementedError( - f"Teacher forcing ignores local_attn_size={local_attn_size}: " - "unlike the block-wise causal mask, this mask always attends " - "to the full clean context. Windowed teacher forcing is not " - "implemented; use local_attn_size=-1 for teacher-forcing " - "training.") + raise NotImplementedError(f"Teacher forcing ignores local_attn_size={local_attn_size}: " + "unlike the block-wise causal mask, this mask always attends " + "to the full clean context. Windowed teacher forcing is not " + "implemented; use local_attn_size=-1 for teacher-forcing " + "training.") total_length = num_frames * frame_seqlen * 2 padded_length = math.ceil(total_length / 128) * 128 - total_length @@ -529,17 +501,19 @@ class CausalWanTransformer3DModel(BaseDiT): noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) attention_block_size = frame_seqlen * num_frame_per_block - frame_indices = torch.arange( - start=0, end=num_frames * frame_seqlen, - step=attention_block_size, device=device, dtype=torch.long - ) + frame_indices = torch.arange(start=0, + end=num_frames * frame_seqlen, + step=attention_block_size, + device=device, + dtype=torch.long) for start in frame_indices: context_ends[start:start + attention_block_size] = start + attention_block_size - noisy_image_start_list = torch.arange( - num_frames * frame_seqlen, total_length, - step=attention_block_size, device=device, dtype=torch.long - ) + noisy_image_start_list = torch.arange(num_frames * frame_seqlen, + total_length, + step=attention_block_size, + device=device, + dtype=torch.long) noisy_image_end_list = noisy_image_start_list + attention_block_size for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)): noise_noise_starts[start:end] = start @@ -554,10 +528,13 @@ class CausalWanTransformer3DModel(BaseDiT): eye_mask = q_idx == kv_idx return eye_mask | clean_mask | noise_mask - block_mask = create_block_mask( - attention_mask, B=None, H=None, - Q_LEN=total_length + padded_length, KV_LEN=total_length + padded_length, - _compile=False, device=device) + block_mask = create_block_mask(attention_mask, + B=None, + H=None, + Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, + _compile=False, + device=device) if not dist.is_initialized() or dist.get_rank() == 0: print(f" cache a teacher-forcing mask with block size of {num_frame_per_block} frames") @@ -565,19 +542,18 @@ class CausalWanTransformer3DModel(BaseDiT): return block_mask - def _forward_inference( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | list[torch.Tensor], - timestep: torch.LongTensor, - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] - | None = None, - kv_cache: dict = None, - crossattn_cache: dict = None, - current_start: int = 0, - cache_start: int = 0, - start_frame: int = 0, - **kwargs) -> torch.Tensor: + def _forward_inference(self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] + | None = None, + kv_cache: dict = None, + crossattn_cache: dict = None, + current_start: int = 0, + cache_start: int = 0, + start_frame: int = 0, + **kwargs) -> torch.Tensor: r""" Run the diffusion model with kv caching. See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. @@ -603,46 +579,44 @@ class CausalWanTransformer3DModel(BaseDiT): rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] if self.rope_cache_policy == "relativistic": # fixed table over [0, max_attention_frames); attention slices it per step - max_attention_frames = ( - GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES - if self.local_attn_size == -1 else self.local_attn_size) + max_attention_frames = (GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES + if self.local_attn_size == -1 else self.local_attn_size) rope_num_frames = max_attention_frames * get_sp_world_size() rope_start_frame = 0 else: rope_num_frames = post_patch_num_frames * get_sp_world_size() rope_start_frame = start_frame # 0 when kv_cache is None - freqs_cos, freqs_sin = get_rotary_pos_embed( - (rope_num_frames, post_patch_height, post_patch_width), - self.hidden_size, - self.num_attention_heads, - rope_dim_list, - dtype=torch.float32 if current_platform.is_mps() else torch.float64, - rope_theta=10000, - start_frame=rope_start_frame - ) + freqs_cos, freqs_sin = get_rotary_pos_embed((rope_num_frames, post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000, + start_frame=rope_start_frame) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = (freqs_cos, - freqs_sin) if freqs_cos is not None else None + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None hidden_states = self.patch_embedding(hidden_states) - grid_sizes = torch.stack( - [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) + grid_sizes = torch.stack([torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) hidden_states = hidden_states.flatten(2).transpose(1, 2) - encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1) + encoder_hidden_states = torch.cat([ + encoder_hidden_states, + encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), + encoder_hidden_states.size(2)) + ], + dim=1) temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) + timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) if encoder_hidden_states_image is not None: - encoder_hidden_states = torch.concat( - [encoder_hidden_states_image, encoder_hidden_states], dim=1) + encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) encoder_hidden_states = encoder_hidden_states.to( - orig_dtype) if current_platform.is_mps( - ) else encoder_hidden_states # cast to orig_dtype for MPS + orig_dtype) if current_platform.is_mps() else encoder_hidden_states # cast to orig_dtype for MPS assert encoder_hidden_states.dtype == orig_dtype @@ -656,28 +630,22 @@ class CausalWanTransformer3DModel(BaseDiT): "block_mask": self.block_mask, "frame_seqlen": post_patch_height * post_patch_width, } - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, - **causal_kwargs) + hidden_states = self._gradient_checkpointing_func(block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, **causal_kwargs) else: causal_kwargs = { "kv_cache": kv_cache[block_index], - "crossattn_cache": (crossattn_cache[block_index] - if crossattn_cache is not None else None), + "crossattn_cache": (crossattn_cache[block_index] if crossattn_cache is not None else None), "current_start": current_start, "cache_start": cache_start, "block_mask": self.block_mask, "frame_seqlen": post_patch_height * post_patch_width, } - hidden_states = block(hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, - **causal_kwargs) + hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, **causal_kwargs) # 5. Output norm, projection & unpatchify temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) - shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, - dim=2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) @@ -686,15 +654,15 @@ class CausalWanTransformer3DModel(BaseDiT): return torch.stack(output) def _forward_train(self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | list[torch.Tensor], - timestep: torch.LongTensor, - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] - | None = None, - start_frame: int = 0, - clean_x: torch.Tensor | None = None, - aug_t: torch.Tensor | None = None, - **kwargs) -> torch.Tensor: + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] + | None = None, + start_frame: int = 0, + clean_x: torch.Tensor | None = None, + aug_t: torch.Tensor | None = None, + **kwargs) -> torch.Tensor: orig_dtype = hidden_states.dtype teacher_forcing = clean_x is not None @@ -714,19 +682,16 @@ class CausalWanTransformer3DModel(BaseDiT): d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames * get_sp_world_size(), post_patch_height, - post_patch_width), + (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), self.hidden_size, self.num_attention_heads, rope_dim_list, dtype=torch.float32 if current_platform.is_mps() else torch.float64, rope_theta=10000, - start_frame=start_frame - ) + start_frame=start_frame) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = (freqs_cos, - freqs_sin) if freqs_cos is not None else None + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None if teacher_forcing: if self.teacher_forcing_block_mask is None: @@ -740,34 +705,35 @@ class CausalWanTransformer3DModel(BaseDiT): block_mask = self.teacher_forcing_block_mask else: if self.block_mask is None: - self.block_mask = self._prepare_blockwise_causal_attn_mask( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=post_patch_height * post_patch_width, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size - ) + self.block_mask = self._prepare_blockwise_causal_attn_mask(device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * + post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size) block_mask = self.block_mask hidden_states = self.patch_embedding(hidden_states) - grid_sizes = torch.stack( - [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) + grid_sizes = torch.stack([torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) hidden_states = hidden_states.flatten(2).transpose(1, 2) - encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1) + encoder_hidden_states = torch.cat([ + encoder_hidden_states, + encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), + encoder_hidden_states.size(2)) + ], + dim=1) encoder_hidden_states_text = encoder_hidden_states temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) + timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) if encoder_hidden_states_image is not None: - encoder_hidden_states = torch.concat( - [encoder_hidden_states_image, encoder_hidden_states], dim=1) + encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) encoder_hidden_states = encoder_hidden_states.to( - orig_dtype) if current_platform.is_mps( - ) else encoder_hidden_states # cast to orig_dtype for MPS + orig_dtype) if current_platform.is_mps() else encoder_hidden_states # cast to orig_dtype for MPS assert encoder_hidden_states.dtype == orig_dtype @@ -777,34 +743,36 @@ class CausalWanTransformer3DModel(BaseDiT): hidden_states = torch.cat([clean_tokens, hidden_states], dim=1) if aug_t is None: aug_t = torch.zeros_like(timestep) - _, timestep_proj_clean, _, _ = self.condition_embedder( - aug_t.flatten(), encoder_hidden_states_text, None) - timestep_proj_clean = timestep_proj_clean.unflatten( - 1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) + _, timestep_proj_clean, _, _ = self.condition_embedder(aug_t.flatten(), encoder_hidden_states_text, None) + timestep_proj_clean = timestep_proj_clean.unflatten(1, + (6, self.hidden_size)).unflatten(dim=0, + sizes=timestep.shape) timestep_proj = torch.cat([timestep_proj_clean, timestep_proj], dim=1) - freqs_cis = (torch.cat([freqs_cos, freqs_cos], dim=0), - torch.cat([freqs_sin, freqs_sin], dim=0)) + freqs_cis = (torch.cat([freqs_cos, freqs_cos], dim=0), torch.cat([freqs_sin, freqs_sin], dim=0)) # 4. Transformer blocks if torch.is_grad_enabled() and self.gradient_checkpointing: for block in self.blocks: - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, - block_mask=block_mask) + hidden_states = self._gradient_checkpointing_func(block, + hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, + block_mask=block_mask) else: for block in self.blocks: - hidden_states = block(hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, - block_mask=block_mask) + hidden_states = block(hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, + block_mask=block_mask) if teacher_forcing: hidden_states = hidden_states[:, hidden_states.shape[1] // 2:] # 5. Output norm, projection & unpatchify temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) - shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, - dim=2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) @@ -812,17 +780,12 @@ class CausalWanTransformer3DModel(BaseDiT): return torch.stack(output) - def forward( - self, - *args, - **kwargs - ): + def forward(self, *args, **kwargs): if kwargs.get('kv_cache', None) is not None: return self._forward_inference(*args, **kwargs) else: return self._forward_train(*args, **kwargs) - def unpatchify(self, x, grid_sizes): r""" @@ -849,5 +812,6 @@ class CausalWanTransformer3DModel(BaseDiT): out.append(u) return out + # Entry point for model registry EntryClass = CausalWanTransformer3DModel diff --git a/fastvideo/models/dits/cosmos.py b/fastvideo/models/dits/cosmos.py index 17b911fdd..d1cf64a4d 100644 --- a/fastvideo/models/dits/cosmos.py +++ b/fastvideo/models/dits/cosmos.py @@ -25,20 +25,14 @@ class CosmosPatchEmbed(nn.Module): super().__init__() self.patch_size = patch_size - self.proj = nn.Linear(in_channels * patch_size[0] * patch_size[1] * - patch_size[2], - out_channels, - bias=bias) + self.proj = nn.Linear(in_channels * patch_size[0] * patch_size[1] * patch_size[2], out_channels, bias=bias) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.patch_size - hidden_states = hidden_states.reshape(batch_size, num_channels, - num_frames // p_t, p_t, - height // p_h, p_h, width // p_w, - p_w) - hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, - 7).flatten(4, 7) + hidden_states = hidden_states.reshape(batch_size, num_channels, num_frames // p_t, p_t, height // p_h, p_h, + width // p_w, p_w) + hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7) hidden_states = self.proj(hidden_states) return hidden_states @@ -63,14 +57,11 @@ class CosmosEmbedding(nn.Module): def __init__(self, embedding_dim: int, condition_dim: int) -> None: super().__init__() - self.time_proj = Timesteps(embedding_dim, - flip_sin_to_cos=True, - downscale_freq_shift=0.0) + self.time_proj = Timesteps(embedding_dim, flip_sin_to_cos=True, downscale_freq_shift=0.0) self.t_embedder = CosmosTimestepEmbedding(embedding_dim, condition_dim) self.norm = RMSNorm(embedding_dim, eps=1e-6) - def forward(self, hidden_states: torch.Tensor, - timestep: torch.LongTensor) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, timestep: torch.LongTensor) -> torch.Tensor: timesteps_proj = self.time_proj(timestep).type_as(hidden_states) temb = self.t_embedder(timesteps_proj) embedded_timestep = self.norm(timesteps_proj) @@ -84,9 +75,7 @@ class CosmosAdaLayerNorm(nn.Module): self.embedding_dim = in_features self.activation = nn.SiLU() - self.norm = nn.LayerNorm(in_features, - elementwise_affine=False, - eps=1e-6) + self.norm = nn.LayerNorm(in_features, elementwise_affine=False, eps=1e-6) self.linear_1 = nn.Linear(in_features, hidden_features, bias=False) self.linear_2 = nn.Linear(hidden_features, 2 * in_features, bias=False) @@ -99,8 +88,7 @@ class CosmosAdaLayerNorm(nn.Module): embedded_timestep = self.linear_2(embedded_timestep) if temb is not None: - embedded_timestep = embedded_timestep + temb[..., :2 * - self.embedding_dim] + embedded_timestep = embedded_timestep + temb[..., :2 * self.embedding_dim] shift, scale = embedded_timestep.chunk(2, dim=-1) with torch.autocast(device_type="cuda", enabled=False): @@ -115,14 +103,10 @@ class CosmosAdaLayerNorm(nn.Module): class CosmosAdaLayerNormZero(nn.Module): - def __init__(self, - in_features: int, - hidden_features: int | None = None) -> None: + def __init__(self, in_features: int, hidden_features: int | None = None) -> None: super().__init__() - self.norm = nn.LayerNorm(in_features, - elementwise_affine=False, - eps=1e-6) + self.norm = nn.LayerNorm(in_features, elementwise_affine=False, eps=1e-6) self.activation = nn.SiLU() if hidden_features is None: @@ -159,12 +143,7 @@ class CosmosAdaLayerNormZero(nn.Module): class CosmosSelfAttention(nn.Module): - def __init__(self, - dim: int, - num_heads: int, - qk_norm=True, - eps=1e-6, - prefix: str = "") -> None: + def __init__(self, dim: int, num_heads: int, qk_norm=True, eps=1e-6, prefix: str = "") -> None: assert dim % num_heads == 0 super().__init__() self.dim = dim @@ -179,11 +158,9 @@ class CosmosSelfAttention(nn.Module): self.to_v = nn.Linear(dim, dim, bias=False) self.to_out = nn.Linear(dim, dim, bias=False) self.dropout = nn.Dropout(0.0) - - self.norm_q = RMSNorm(self.head_dim, - eps=eps) if qk_norm else nn.Identity() - self.norm_k = RMSNorm(self.head_dim, - eps=eps) if qk_norm else nn.Identity() + + self.norm_q = RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity() def forward(self, hidden_states: torch.Tensor, @@ -212,14 +189,8 @@ class CosmosSelfAttention(nn.Module): # Apply RoPE if provided if image_rotary_emb is not None: - query = apply_rotary_emb(query, - image_rotary_emb, - use_real=True, - use_real_unbind_dim=-2) - key = apply_rotary_emb(key, - image_rotary_emb, - use_real=True, - use_real_unbind_dim=-2) + query = apply_rotary_emb(query, image_rotary_emb, use_real=True, use_real_unbind_dim=-2) + key = apply_rotary_emb(key, image_rotary_emb, use_real=True, use_real_unbind_dim=-2) # Prepare for GQA (Grouped Query Attention) if torch.onnx.is_in_onnx_export(): @@ -234,9 +205,12 @@ class CosmosSelfAttention(nn.Module): value = value.repeat_interleave(query_idx // value_idx, dim=3) # Attention computation - attn_output = torch.nn.functional.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) + attn_output = torch.nn.functional.scaled_dot_product_attention(query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False) attn_output = attn_output.transpose(1, 2).flatten(2, 3).type_as(query) # Output projection @@ -269,11 +243,9 @@ class CosmosCrossAttention(nn.Module): self.to_v = nn.Linear(cross_attention_dim, dim, bias=False) self.to_out = nn.Linear(dim, dim, bias=False) self.dropout = nn.Dropout(0.0) - - self.norm_q = RMSNorm(self.head_dim, - eps=eps) if qk_norm else nn.Identity() - self.norm_k = RMSNorm(self.head_dim, - eps=eps) if qk_norm else nn.Identity() + + self.norm_q = RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity() def forward(self, hidden_states: torch.Tensor, @@ -309,9 +281,12 @@ class CosmosCrossAttention(nn.Module): value = value.repeat_interleave(query_idx // value_idx, dim=3) # Attention computation - attn_output = torch.nn.functional.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) + attn_output = torch.nn.functional.scaled_dot_product_attention(query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False) attn_output = attn_output.transpose(1, 2).flatten(2, 3).type_as(query) # Output projection @@ -338,31 +313,23 @@ class CosmosTransformerBlock(nn.Module): hidden_size = num_attention_heads * attention_head_dim - self.norm1 = CosmosAdaLayerNormZero(in_features=hidden_size, - hidden_features=adaln_lora_dim) - self.attn1 = CosmosSelfAttention( - dim=hidden_size, - num_heads=num_attention_heads, - qk_norm=(qk_norm == "rms_norm"), - eps=1e-5, - prefix=f"{prefix}.attn1") + self.norm1 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim) + self.attn1 = CosmosSelfAttention(dim=hidden_size, + num_heads=num_attention_heads, + qk_norm=(qk_norm == "rms_norm"), + eps=1e-5, + prefix=f"{prefix}.attn1") - self.norm2 = CosmosAdaLayerNormZero(in_features=hidden_size, - hidden_features=adaln_lora_dim) - self.attn2 = CosmosCrossAttention( - dim=hidden_size, - cross_attention_dim=cross_attention_dim, - num_heads=num_attention_heads, - qk_norm=(qk_norm == "rms_norm"), - eps=1e-5, - prefix=f"{prefix}.attn2") + self.norm2 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim) + self.attn2 = CosmosCrossAttention(dim=hidden_size, + cross_attention_dim=cross_attention_dim, + num_heads=num_attention_heads, + qk_norm=(qk_norm == "rms_norm"), + eps=1e-5, + prefix=f"{prefix}.attn2") - self.norm3 = CosmosAdaLayerNormZero(in_features=hidden_size, - hidden_features=adaln_lora_dim) - self.ff = MLP(hidden_size, - int(hidden_size * mlp_ratio), - act_type="gelu", - bias=False) + self.norm3 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim) + self.ff = MLP(hidden_size, int(hidden_size * mlp_ratio), act_type="gelu", bias=False) def forward( self, @@ -377,23 +344,19 @@ class CosmosTransformerBlock(nn.Module): if extra_pos_emb is not None: hidden_states = hidden_states + extra_pos_emb - norm_hidden_states, gate = self.norm1(hidden_states, embedded_timestep, - temb) + norm_hidden_states, gate = self.norm1(hidden_states, embedded_timestep, temb) - attn_output = self.attn1(norm_hidden_states, - image_rotary_emb=image_rotary_emb) + attn_output = self.attn1(norm_hidden_states, image_rotary_emb=image_rotary_emb) hidden_states = hidden_states + gate * attn_output - norm_hidden_states, gate = self.norm2(hidden_states, embedded_timestep, - temb) + norm_hidden_states, gate = self.norm2(hidden_states, embedded_timestep, temb) attn_output = self.attn2(norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask) hidden_states = hidden_states + gate * attn_output - norm_hidden_states, gate = self.norm3(hidden_states, embedded_timestep, - temb) + norm_hidden_states, gate = self.norm3(hidden_states, embedded_timestep, temb) ff_output = self.ff(norm_hidden_states) hidden_states = hidden_states + gate * ff_output @@ -412,10 +375,7 @@ class CosmosRotaryPosEmbed(nn.Module): ) -> None: super().__init__() - self.max_size = [ - size // patch - for size, patch in zip(max_size, patch_size, strict=False) - ] + self.max_size = [size // patch for size, patch in zip(max_size, patch_size, strict=False)] self.patch_size = patch_size self.base_fps = base_fps @@ -427,56 +387,39 @@ class CosmosRotaryPosEmbed(nn.Module): self.w_ntk_factor = rope_scale[2]**(self.dim_w / (self.dim_w - 2)) self.t_ntk_factor = rope_scale[0]**(self.dim_t / (self.dim_t - 2)) - - def forward(self, - hidden_states: torch.Tensor, - fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]: + def forward(self, hidden_states: torch.Tensor, fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]: fps = 16 batch_size, num_channels, num_frames, height, width = hidden_states.shape - pe_size = [ - num_frames // self.patch_size[0], height // self.patch_size[1], - width // self.patch_size[2] - ] + pe_size = [num_frames // self.patch_size[0], height // self.patch_size[1], width // self.patch_size[2]] device = hidden_states.device h_theta = 10000.0 * self.h_ntk_factor w_theta = 10000.0 * self.w_ntk_factor t_theta = 10000.0 * self.t_ntk_factor - seq = torch.arange(max(self.max_size), - device=device, - dtype=torch.float32) - dim_h_range = ( - torch.arange(0, self.dim_h, 2, device=device, - dtype=torch.float32)[:(self.dim_h // 2)] / self.dim_h) - dim_w_range = ( - torch.arange(0, self.dim_w, 2, device=device, - dtype=torch.float32)[:(self.dim_w // 2)] / self.dim_w) - dim_t_range = ( - torch.arange(0, self.dim_t, 2, device=device, - dtype=torch.float32)[:(self.dim_t // 2)] / self.dim_t) + seq = torch.arange(max(self.max_size), device=device, dtype=torch.float32) + dim_h_range = (torch.arange(0, self.dim_h, 2, device=device, dtype=torch.float32)[:(self.dim_h // 2)] / + self.dim_h) + dim_w_range = (torch.arange(0, self.dim_w, 2, device=device, dtype=torch.float32)[:(self.dim_w // 2)] / + self.dim_w) + dim_t_range = (torch.arange(0, self.dim_t, 2, device=device, dtype=torch.float32)[:(self.dim_t // 2)] / + self.dim_t) h_spatial_freqs = 1.0 / (h_theta**dim_h_range) w_spatial_freqs = 1.0 / (w_theta**dim_w_range) temporal_freqs = 1.0 / (t_theta**dim_t_range) - emb_h = torch.outer(seq[:pe_size[1]], - h_spatial_freqs)[None, :, None, :].repeat( - pe_size[0], 1, pe_size[2], 1) - emb_w = torch.outer(seq[:pe_size[2]], - w_spatial_freqs)[None, None, :, :].repeat( - pe_size[0], pe_size[1], 1, 1) + emb_h = torch.outer(seq[:pe_size[1]], h_spatial_freqs)[None, :, None, :].repeat(pe_size[0], 1, pe_size[2], 1) + emb_w = torch.outer(seq[:pe_size[2]], w_spatial_freqs)[None, None, :, :].repeat(pe_size[0], pe_size[1], 1, 1) if fps is None: emb_t = torch.outer(seq[:pe_size[0]], temporal_freqs) else: temporal_scale = seq[:pe_size[0]] / fps * self.base_fps - emb_t = torch.outer(temporal_scale, - temporal_freqs) + emb_t = torch.outer(temporal_scale, temporal_freqs) emb_t = emb_t[:, None, None, :].repeat(1, pe_size[1], pe_size[2], 1) - freqs = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1).flatten(0, - 2).float() + freqs = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1).flatten(0, 2).float() cos = torch.cos(freqs) sin = torch.sin(freqs) return cos, sin @@ -493,43 +436,26 @@ class CosmosLearnablePositionalEmbed(nn.Module): ) -> None: super().__init__() - self.max_size = [ - size // patch - for size, patch in zip(max_size, patch_size, strict=False) - ] + self.max_size = [size // patch for size, patch in zip(max_size, patch_size, strict=False)] self.patch_size = patch_size self.eps = eps - self.pos_emb_t = nn.Parameter(torch.zeros(self.max_size[0], - hidden_size)) - self.pos_emb_h = nn.Parameter(torch.zeros(self.max_size[1], - hidden_size)) - self.pos_emb_w = nn.Parameter(torch.zeros(self.max_size[2], - hidden_size)) + self.pos_emb_t = nn.Parameter(torch.zeros(self.max_size[0], hidden_size)) + self.pos_emb_h = nn.Parameter(torch.zeros(self.max_size[1], hidden_size)) + self.pos_emb_w = nn.Parameter(torch.zeros(self.max_size[2], hidden_size)) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = hidden_states.shape - pe_size = [ - num_frames // self.patch_size[0], height // self.patch_size[1], - width // self.patch_size[2] - ] + pe_size = [num_frames // self.patch_size[0], height // self.patch_size[1], width // self.patch_size[2]] - emb_t = self.pos_emb_t[:pe_size[0]][None, :, None, None, :].repeat( - batch_size, 1, pe_size[1], pe_size[2], 1) - emb_h = self.pos_emb_h[:pe_size[1]][None, None, :, None, :].repeat( - batch_size, pe_size[0], 1, pe_size[2], 1) - emb_w = self.pos_emb_w[:pe_size[2]][None, None, None, :, :].repeat( - batch_size, pe_size[0], pe_size[1], 1, 1) + emb_t = self.pos_emb_t[:pe_size[0]][None, :, None, None, :].repeat(batch_size, 1, pe_size[1], pe_size[2], 1) + emb_h = self.pos_emb_h[:pe_size[1]][None, None, :, None, :].repeat(batch_size, pe_size[0], 1, pe_size[2], 1) + emb_w = self.pos_emb_w[:pe_size[2]][None, None, None, :, :].repeat(batch_size, pe_size[0], pe_size[1], 1, 1) emb = emb_t + emb_h + emb_w emb = emb.flatten(1, 3) - norm = torch.linalg.vector_norm(emb, - dim=-1, - keepdim=True, - dtype=torch.float32) - norm = torch.add(self.eps, - norm, - alpha=np.sqrt(norm.numel() / emb.numel())) + norm = torch.linalg.vector_norm(emb, dim=-1, keepdim=True, dtype=torch.float32) + norm = torch.add(self.eps, norm, alpha=np.sqrt(norm.numel() / emb.numel())) return (emb / norm).type_as(hidden_states) @@ -560,10 +486,7 @@ class CosmosTransformer3DModel(BaseDiT): # (HF config: in_channels=17 = 16 latent + 1 condition_mask). # Only add +1 for the padding_mask when concat_padding_mask=True. patch_embed_in_channels = config.in_channels + (1 if config.concat_padding_mask else 0) - self.patch_embed = CosmosPatchEmbed(patch_embed_in_channels, - inner_dim, - config.patch_size, - bias=False) + self.patch_embed = CosmosPatchEmbed(patch_embed_in_channels, inner_dim, config.patch_size, bias=False) # 2. Positional Embedding self.rope = CosmosRotaryPosEmbed(hidden_size=config.attention_head_dim, @@ -598,10 +521,7 @@ class CosmosTransformer3DModel(BaseDiT): # 5. Output norm & projection self.norm_out = CosmosAdaLayerNorm(inner_dim, config.adaln_lora_dim) - self.proj_out = nn.Linear(inner_dim, - config.out_channels * - math.prod(config.patch_size), - bias=False) + self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size), bias=False) self.gradient_checkpointing = False self.__post_init__() @@ -638,8 +558,7 @@ class CosmosTransformer3DModel(BaseDiT): encoder_hidden_states = encoder_hidden_states.to(_target_dtype) else: encoder_hidden_states = [ - t.to(_target_dtype) if t.dtype != _target_dtype else t - for t in encoder_hidden_states + t.to(_target_dtype) if t.dtype != _target_dtype else t for t in encoder_hidden_states ] # 1. Concatenate padding mask if needed & prepare attention mask @@ -648,25 +567,22 @@ class CosmosTransformer3DModel(BaseDiT): if self.concat_padding_mask: from torchvision import transforms - padding_mask = transforms.functional.resize( - padding_mask, list(hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST - ) + padding_mask = transforms.functional.resize(padding_mask, + list(hidden_states.shape[-2:]), + interpolation=transforms.InterpolationMode.NEAREST) # torchvision.resize may upcast bf16/fp16 → fp32; restore # hidden_states' dtype so the subsequent cat doesn't promote # everything and break patch_embed (bf16 weights). padding_mask = padding_mask.to(hidden_states.dtype) hidden_states = torch.cat( - [hidden_states, padding_mask.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)], dim=1 - ) + [hidden_states, padding_mask.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)], dim=1) if attention_mask is not None: - attention_mask = attention_mask.unsqueeze(1).unsqueeze( - 1) # [B, 1, 1, S] + attention_mask = attention_mask.unsqueeze(1).unsqueeze(1) # [B, 1, 1, S] # 2. Generate positional embeddings image_rotary_emb = self.rope(hidden_states, fps=fps) - extra_pos_emb = self.learnable_pos_embed( - hidden_states) if self.extra_pos_embed_type == "learnable" else None + extra_pos_emb = self.learnable_pos_embed(hidden_states) if self.extra_pos_embed_type == "learnable" else None # 3. Patchify input p_t, p_h, p_w = self.patch_size @@ -674,23 +590,20 @@ class CosmosTransformer3DModel(BaseDiT): post_patch_height = height // p_h post_patch_width = width // p_w hidden_states = self.patch_embed(hidden_states) - hidden_states = hidden_states.flatten( - 1, 3) # [B, T, H, W, C] -> [B, THW, C] codespell:ignore + hidden_states = hidden_states.flatten(1, 3) # [B, T, H, W, C] -> [B, THW, C] codespell:ignore # 4. Timestep embeddings if timestep.ndim == 1: temb, embedded_timestep = self.time_embed(hidden_states, timestep) elif timestep.ndim == 5: - assert timestep.shape == (batch_size, 1, num_frames, 1, 1), ( - f"Expected timestep to have shape [B, 1, T, 1, 1], but got {timestep.shape}" - ) + assert timestep.shape == (batch_size, 1, num_frames, 1, + 1), (f"Expected timestep to have shape [B, 1, T, 1, 1], but got {timestep.shape}") timestep = timestep.flatten() temb, embedded_timestep = self.time_embed(hidden_states, timestep) # We can do this because num_frames == post_patch_num_frames, as p_t is 1 temb, embedded_timestep = ( - x.view(batch_size, post_patch_num_frames, 1, 1, - -1).expand(-1, -1, post_patch_height, post_patch_width, - -1).flatten(1, 3) + x.view(batch_size, post_patch_num_frames, 1, 1, -1).expand(-1, -1, post_patch_height, post_patch_width, + -1).flatten(1, 3) for x in (temb, embedded_timestep) ) # [BT, C] -> [B, T, 1, 1, C] -> [B, T, H, W, C] -> [B, THW, C] codespell:ignore else: @@ -725,8 +638,7 @@ class CosmosTransformer3DModel(BaseDiT): hidden_states = self.norm_out(hidden_states, embedded_timestep, temb) hidden_states = self.proj_out(hidden_states) hidden_states = hidden_states.unflatten(2, (p_h, p_w, p_t, -1)) - hidden_states = hidden_states.unflatten( - 1, (post_patch_num_frames, post_patch_height, post_patch_width)) + hidden_states = hidden_states.unflatten(1, (post_patch_num_frames, post_patch_height, post_patch_width)) # NOTE: The permutation order here is not the inverse operation of what happens when patching as usually expected. # It might be a source of confusion to the reader, but this is correct hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5) @@ -736,8 +648,9 @@ class CosmosTransformer3DModel(BaseDiT): # do `transformer(..., return_dict=False)[0]` (diffusers # convention used by CosmosDenoisingStage). if not kwargs.get("return_dict", True): - return (hidden_states,) + return (hidden_states, ) return hidden_states + # Entry point for model registry EntryClass = CosmosTransformer3DModel diff --git a/fastvideo/models/dits/cosmos2_5.py b/fastvideo/models/dits/cosmos2_5.py index c7e67f065..c35bd88e9 100644 --- a/fastvideo/models/dits/cosmos2_5.py +++ b/fastvideo/models/dits/cosmos2_5.py @@ -20,7 +20,6 @@ from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum - class Cosmos25PatchEmbed(nn.Module): """ COSMOS 2.5 patch embedding - converts video (B, C, T, H, W) to patches (B, T', H', W', D). @@ -50,12 +49,8 @@ class Cosmos25PatchEmbed(nn.Module): p_t, p_h, p_w = self.patch_size # Rearrange: b c (t pt) (h ph) (w pw) -> b t h w (c pt ph pw) - hidden_states = hidden_states.reshape( - batch_size, num_channels, - num_frames // p_t, p_t, - height // p_h, p_h, - width // p_w, p_w - ) + hidden_states = hidden_states.reshape(batch_size, num_channels, num_frames // p_t, p_t, height // p_h, p_h, + width // p_w, p_w) hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7) hidden_states = hidden_states.flatten(4, 7) # Flatten patch dimensions @@ -80,10 +75,10 @@ class Cosmos25TimestepEmbedding(nn.Module): ) -> None: super().__init__() self.use_adaln_lora = use_adaln_lora - + self.linear_1 = nn.Linear(in_features, out_features, bias=False) self.activation = nn.SiLU() - + if use_adaln_lora: self.linear_2 = nn.Linear(out_features, 3 * out_features, bias=False) else: @@ -149,15 +144,15 @@ class Cosmos25Embedding(nn.Module): # Handle 2D timestep input (B, T) like the official model assert timestep.ndim == 2, f"Expected 2D timestep, got {timestep.ndim}D with shape {timestep.shape}" B, T = timestep.shape - + # Flatten for Timesteps layer which expects 1D, then reshape back timestep_flat = timestep.flatten() # (B*T,) timesteps_proj = self.time_proj(timestep_flat).type_as(hidden_states) # (B*T, D) timesteps_proj = timesteps_proj.reshape(B, T, -1) # (B, T, D) - + embedded_timestep, adaln_lora = self.t_embedder(timesteps_proj) embedded_timestep = self.norm(embedded_timestep) - + return embedded_timestep, adaln_lora @@ -226,15 +221,13 @@ class Cosmos25SelfAttention(nn.Module): # For single-GPU (non-distributed), use LocalAttention to avoid distributed requirements if supported_attention_backends is None: supported_attention_backends = (AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA) - + # Always use DistributedAttention (requires distributed environment to be initialized) - self.attn = DistributedAttention( - num_heads=num_heads, - head_size=self.head_dim, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix="self_attn" - ) + self.attn = DistributedAttention(num_heads=num_heads, + head_size=self.head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix="self_attn") def forward( self, @@ -271,8 +264,7 @@ class Cosmos25SelfAttention(nn.Module): query = query.transpose(1, 2) # (B, H, S, D_h) -> (B, S, H, D_h) key = key.transpose(1, 2) value = value.transpose(1, 2) - - + attn_output, _ = self.attn(query, key, value) # Reshape back: (B, S, H, D_h) -> (B, S, H*D_h) attn_output = attn_output.flatten(-2, -1) @@ -313,7 +305,7 @@ class Cosmos25CrossAttention(nn.Module): if supported_attention_backends is None: supported_attention_backends = (AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA) - + # Use LocalAttention for cross-attention since text embeddings are not sharded # in sequence parallelism (replicated across ranks) self.attn = LocalAttention( @@ -422,15 +414,11 @@ class Cosmos25TransformerBlock(nn.Module): nn.Linear(adaln_lora_dim, 3 * hidden_size, bias=False), ) else: - self.adaln_modulation_self_attn = nn.Sequential( - nn.SiLU(), nn.Linear(hidden_size, 3 * hidden_size, bias=False) - ) - self.adaln_modulation_cross_attn = nn.Sequential( - nn.SiLU(), nn.Linear(hidden_size, 3 * hidden_size, bias=False) - ) - self.adaln_modulation_mlp = nn.Sequential( - nn.SiLU(), nn.Linear(hidden_size, 3 * hidden_size, bias=False) - ) + self.adaln_modulation_self_attn = nn.Sequential(nn.SiLU(), + nn.Linear(hidden_size, 3 * hidden_size, bias=False)) + self.adaln_modulation_cross_attn = nn.Sequential(nn.SiLU(), + nn.Linear(hidden_size, 3 * hidden_size, bias=False)) + self.adaln_modulation_mlp = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 3 * hidden_size, bias=False)) def forward( self, @@ -459,22 +447,17 @@ class Cosmos25TransformerBlock(nn.Module): # Step 1: Compute ALL modulation parameters once (matches official model) if self.use_adaln_lora and adaln_lora is not None: - shift_self_attn, scale_self_attn, gate_self_attn = ( - self.adaln_modulation_self_attn(embedded_timestep) + adaln_lora - ).chunk(3, dim=-1) - shift_cross_attn, scale_cross_attn, gate_cross_attn = ( - self.adaln_modulation_cross_attn(embedded_timestep) + adaln_lora - ).chunk(3, dim=-1) - shift_mlp, scale_mlp, gate_mlp = ( - self.adaln_modulation_mlp(embedded_timestep) + adaln_lora - ).chunk(3, dim=-1) + shift_self_attn, scale_self_attn, gate_self_attn = (self.adaln_modulation_self_attn(embedded_timestep) + + adaln_lora).chunk(3, dim=-1) + shift_cross_attn, scale_cross_attn, gate_cross_attn = (self.adaln_modulation_cross_attn(embedded_timestep) + + adaln_lora).chunk(3, dim=-1) + shift_mlp, scale_mlp, gate_mlp = (self.adaln_modulation_mlp(embedded_timestep) + adaln_lora).chunk(3, + dim=-1) else: - shift_self_attn, scale_self_attn, gate_self_attn = self.adaln_modulation_self_attn( - embedded_timestep - ).chunk(3, dim=-1) + shift_self_attn, scale_self_attn, gate_self_attn = self.adaln_modulation_self_attn(embedded_timestep).chunk( + 3, dim=-1) shift_cross_attn, scale_cross_attn, gate_cross_attn = self.adaln_modulation_cross_attn( - embedded_timestep - ).chunk(3, dim=-1) + embedded_timestep).chunk(3, dim=-1) shift_mlp, scale_mlp, gate_mlp = self.adaln_modulation_mlp(embedded_timestep).chunk(3, dim=-1) # Reshape modulation parameters from (B, T, D) to (B, T, 1, 1, D) for broadcasting @@ -494,9 +477,9 @@ class Cosmos25TransformerBlock(nn.Module): norm_hidden_states = self.norm1(hidden_states, shift_self_attn, scale_self_attn) # Flatten for attention: (B, T, H, W, D) -> (B, THW, D) norm_hidden_states_flat = norm_hidden_states.flatten(1, 3) - + attn_output = self.attn1(norm_hidden_states_flat, rope_emb=rope_emb) - + # Reshape back and apply residual attn_output = attn_output.unflatten(1, (T, H, W)) # (B, T, H, W, D) hidden_states = hidden_states + gate_self_attn * attn_output @@ -504,22 +487,22 @@ class Cosmos25TransformerBlock(nn.Module): # Step 3: Cross-attention block norm_hidden_states = self.norm2(hidden_states, shift_cross_attn, scale_cross_attn) norm_hidden_states_flat = norm_hidden_states.flatten(1, 3) - + attn_output = self.attn2( norm_hidden_states_flat, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask, ) - + attn_output = attn_output.unflatten(1, (T, H, W)) hidden_states = hidden_states + gate_cross_attn * attn_output # Step 4: MLP block norm_hidden_states = self.norm3(hidden_states, shift_mlp, scale_mlp) norm_hidden_states_flat = norm_hidden_states.flatten(1, 3) - + mlp_output = self.mlp(norm_hidden_states_flat) - + mlp_output = mlp_output.unflatten(1, (T, H, W)) hidden_states = hidden_states + gate_mlp * mlp_output @@ -532,13 +515,13 @@ class Cosmos25RotaryPosEmbed(nn.Module): """ def __init__( - self, - hidden_size: int, - max_size: tuple[int, int, int] = (128, 240, 240), - patch_size: tuple[int, int, int] = (1, 2, 2), - base_fps: int = 24, - rope_scale: tuple[float, float, float] = (1.0, 1.0, 1.0), - enable_fps_modulation: bool = True, + self, + hidden_size: int, + max_size: tuple[int, int, int] = (128, 240, 240), + patch_size: tuple[int, int, int] = (1, 2, 2), + base_fps: int = 24, + rope_scale: tuple[float, float, float] = (1.0, 1.0, 1.0), + enable_fps_modulation: bool = True, ) -> None: super().__init__() @@ -553,13 +536,11 @@ class Cosmos25RotaryPosEmbed(nn.Module): self.dim_t = hidden_size - self.dim_h - self.dim_w # NTK-aware extrapolation factors - self.h_ntk_factor = rope_scale[1] ** (self.dim_h / (self.dim_h - 2)) - self.w_ntk_factor = rope_scale[2] ** (self.dim_w / (self.dim_w - 2)) - self.t_ntk_factor = rope_scale[0] ** (self.dim_t / (self.dim_t - 2)) + self.h_ntk_factor = rope_scale[1]**(self.dim_h / (self.dim_h - 2)) + self.w_ntk_factor = rope_scale[2]**(self.dim_w / (self.dim_w - 2)) + self.t_ntk_factor = rope_scale[0]**(self.dim_t / (self.dim_t - 2)) - def forward( - self, hidden_states: torch.Tensor, fps: int | None = None - ) -> tuple[torch.Tensor, torch.Tensor]: + def forward(self, hidden_states: torch.Tensor, fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]: """ Generate 3D RoPE embeddings. @@ -584,13 +565,16 @@ class Cosmos25RotaryPosEmbed(nn.Module): seq = torch.arange(max(self.max_size), device=device, dtype=torch.float32) # Use self.dim_h/w/t which were set during initialization - dim_h_range = torch.arange(0, self.dim_h, 2, device=device, dtype=torch.float32)[: (self.dim_h // 2)] / self.dim_h - dim_w_range = torch.arange(0, self.dim_w, 2, device=device, dtype=torch.float32)[: (self.dim_w // 2)] / self.dim_w - dim_t_range = torch.arange(0, self.dim_t, 2, device=device, dtype=torch.float32)[: (self.dim_t // 2)] / self.dim_t + dim_h_range = torch.arange(0, self.dim_h, 2, device=device, + dtype=torch.float32)[:(self.dim_h // 2)] / self.dim_h + dim_w_range = torch.arange(0, self.dim_w, 2, device=device, + dtype=torch.float32)[:(self.dim_w // 2)] / self.dim_w + dim_t_range = torch.arange(0, self.dim_t, 2, device=device, + dtype=torch.float32)[:(self.dim_t // 2)] / self.dim_t - h_spatial_freqs = 1.0 / (h_theta ** dim_h_range) - w_spatial_freqs = 1.0 / (w_theta ** dim_w_range) - temporal_freqs = 1.0 / (t_theta ** dim_t_range) + h_spatial_freqs = 1.0 / (h_theta**dim_h_range) + w_spatial_freqs = 1.0 / (w_theta**dim_w_range) + temporal_freqs = 1.0 / (t_theta**dim_t_range) # Generate positional embeddings half_emb_h = torch.outer(seq[:H], h_spatial_freqs) @@ -613,7 +597,7 @@ class Cosmos25RotaryPosEmbed(nn.Module): cos = torch.cos(freqs) # (THW, D) sin = torch.sin(freqs) # (THW, D) - + return cos, sin @@ -710,7 +694,7 @@ class Cosmos25FinalLayer(nn.Module): if self.use_adaln_lora and adaln_lora is not None: # Use first 2*hidden_size elements for shift/scale - embedded_timestep = embedded_timestep + adaln_lora[..., : 2 * self.hidden_size] + embedded_timestep = embedded_timestep + adaln_lora[..., :2 * self.hidden_size] shift, scale = embedded_timestep.chunk(2, dim=-1) @@ -742,7 +726,7 @@ class Cosmos25Transformer3DModel(BaseDiT): - QK normalization - Cross-attention projection (optional) """ - + _fsdp_shard_conditions = Cosmos25VideoConfig()._fsdp_shard_conditions _compile_conditions = Cosmos25VideoConfig()._compile_conditions param_names_mapping = Cosmos25VideoConfig().param_names_mapping @@ -773,10 +757,8 @@ class Cosmos25Transformer3DModel(BaseDiT): if config.concat_padding_mask: patch_embed_in_channels += 1 # Add 1 for padding_mask # Total: 16 + 1 + 1 = 18 (with concat_padding_mask=True) - - self.patch_embed = Cosmos25PatchEmbed( - patch_embed_in_channels, inner_dim, config.patch_size - ) + + self.patch_embed = Cosmos25PatchEmbed(patch_embed_in_channels, inner_dim, config.patch_size) # 2. Positional Embeddings self.rope = Cosmos25RotaryPosEmbed( @@ -822,8 +804,7 @@ class Cosmos25Transformer3DModel(BaseDiT): use_adaln_lora=self.use_adaln_lora, qk_norm=(config.qk_norm == "rms_norm"), supported_attention_backends=config._supported_attention_backends, - ) - for i in range(config.num_layers) + ) for i in range(config.num_layers) ]) # 6. Final Layer @@ -885,10 +866,8 @@ class Cosmos25Transformer3DModel(BaseDiT): post_patch_num_frames = num_frames // p_t post_patch_height = height // p_h post_patch_width = width // p_w - - + hidden_states = self.patch_embed(hidden_states) # (B, T', H', W', D) - # 4. Generate RoPE embeddings (after patchify, using patch dimensions) rope_emb = self.rope(hidden_states, fps=fps) @@ -908,18 +887,13 @@ class Cosmos25Transformer3DModel(BaseDiT): pass else: raise ValueError(f"Unsupported timestep shape: {timestep.shape}") - + # Now timestep is always (B, T), pass directly to time_embed embedded_timestep, adaln_lora = self.time_embed(hidden_states, timestep) # 7. Apply cross-attention projection (if used) if self.use_crossattn_projection: - encoder_hidden_states = self.crossattn_proj( - encoder_hidden_states.to(self.crossattn_proj[0].weight.dtype) - ) - - - + encoder_hidden_states = self.crossattn_proj(encoder_hidden_states.to(self.crossattn_proj[0].weight.dtype)) # Prepare attention mask if attention_mask is not None: @@ -963,5 +937,6 @@ class Cosmos25Transformer3DModel(BaseDiT): return hidden_states + # Entry point for model registry EntryClass = Cosmos25Transformer3DModel diff --git a/fastvideo/models/dits/dreamx_world.py b/fastvideo/models/dits/dreamx_world.py index e80164d6b..a53f50ec3 100644 --- a/fastvideo/models/dits/dreamx_world.py +++ b/fastvideo/models/dits/dreamx_world.py @@ -7,18 +7,14 @@ import torch import torch.nn as nn from fastvideo.configs.models.dits.dreamx_world import DreamXWorldConfig -from fastvideo.distributed.communication_op import ( - sequence_model_parallel_all_gather_with_unpad, - sequence_model_parallel_shard) +from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad, + sequence_model_parallel_shard) from fastvideo.distributed.parallel_state import get_sp_world_size from fastvideo.layers.layernorm import RMSNorm from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.rotary_embedding import get_rotary_pos_embed -from fastvideo.models.dits.wanvideo import (LayerNormScaleShift, - PatchEmbed, - WanTimeTextImageEmbedding, - WanTransformer3DModel, - WanTransformerBlock) +from fastvideo.models.dits.wanvideo import (LayerNormScaleShift, PatchEmbed, WanTimeTextImageEmbedding, + WanTransformer3DModel, WanTransformerBlock) from fastvideo.platforms import AttentionBackendEnum, current_platform from fastvideo.attention import LocalAttention from fastvideo.distributed.parallel_state import get_sp_world_size @@ -31,17 +27,14 @@ def _dreamx_invert_se3(transforms: torch.Tensor) -> torch.Tensor: rot_inv = transforms[..., :3, :3].transpose(-1, -2) out = torch.zeros_like(transforms) out[..., :3, :3] = rot_inv - out[..., :3, 3] = -torch.einsum("...ij,...j->...i", rot_inv, - transforms[..., :3, 3]) + out[..., :3, 3] = -torch.einsum("...ij,...j->...i", rot_inv, transforms[..., :3, 3]) out[..., 3, 3] = 1.0 return out.to(dtype=transforms.dtype) def _dreamx_lift_k(intrinsics: torch.Tensor) -> torch.Tensor: assert intrinsics.shape[-2:] == (3, 3) - out = torch.zeros(intrinsics.shape[:-2] + (4, 4), - device=intrinsics.device, - dtype=intrinsics.dtype) + out = torch.zeros(intrinsics.shape[:-2] + (4, 4), device=intrinsics.device, dtype=intrinsics.dtype) out[..., :3, :3] = intrinsics out[..., 3, 3] = 1.0 return out @@ -58,28 +51,25 @@ def _dreamx_invert_k(intrinsics: torch.Tensor) -> torch.Tensor: return out.to(dtype=intrinsics.dtype) -def _dreamx_apply_tiled_projmat(feats: torch.Tensor, - matrix: torch.Tensor) -> torch.Tensor: +def _dreamx_apply_tiled_projmat(feats: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor: batch, num_heads, seq_len, feat_dim = feats.shape proj_dim = matrix.shape[-1] assert feat_dim % proj_dim == 0 if matrix.shape[1] == seq_len: - feats = feats.view(batch, num_heads, seq_len, feat_dim // proj_dim, - proj_dim) + feats = feats.view(batch, num_heads, seq_len, feat_dim // proj_dim, proj_dim) out = torch.einsum("btij,bntpj->bntpi", matrix, feats) return out.reshape(batch, num_heads, seq_len, feat_dim) cameras = matrix.shape[1] assert seq_len > cameras and seq_len % cameras == 0 - feats = feats.reshape(batch, num_heads, cameras, -1, - feat_dim // proj_dim, proj_dim) + feats = feats.reshape(batch, num_heads, cameras, -1, feat_dim // proj_dim, proj_dim) out = torch.einsum("bcij,bncpkj->bncpki", matrix, feats) return out.reshape(batch, num_heads, seq_len, feat_dim) -def _dreamx_prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - viewmats: torch.Tensor, intrinsics: torch.Tensor): +def _dreamx_prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, viewmats: torch.Tensor, + intrinsics: torch.Tensor): batch, num_heads, seq_len, head_dim = q.shape cameras = viewmats.shape[1] assert q.shape == k.shape == v.shape @@ -92,8 +82,7 @@ def _dreamx_prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, intrinsics_norm[..., 1, 1] = intrinsics[..., 1, 1] intrinsics_norm[..., 2, 2] = 1.0 - proj = torch.einsum("...ij,...jk->...ik", - _dreamx_lift_k(intrinsics_norm), viewmats) + proj = torch.einsum("...ij,...jk->...ik", _dreamx_lift_k(intrinsics_norm), viewmats) proj_t = proj.transpose(-1, -2).to(dtype=viewmats.dtype) proj_inv = torch.einsum( "...ij,...jk->...ik", @@ -125,22 +114,10 @@ class DreamXPropeSelfAttention(nn.Module): self.head_dim = attn_dim // num_heads self.qk_norm = qk_norm - self.q_proj = ReplicatedLinear(dim, - attn_dim, - quant_config=quant_config, - prefix=f"{prefix}.q_proj") - self.k_proj = ReplicatedLinear(dim, - attn_dim, - quant_config=quant_config, - prefix=f"{prefix}.k_proj") - self.v_proj = ReplicatedLinear(dim, - attn_dim, - quant_config=quant_config, - prefix=f"{prefix}.v_proj") - self.out_proj = ReplicatedLinear(attn_dim, - dim, - quant_config=quant_config, - prefix=f"{prefix}.out_proj") + self.q_proj = ReplicatedLinear(dim, attn_dim, quant_config=quant_config, prefix=f"{prefix}.q_proj") + self.k_proj = ReplicatedLinear(dim, attn_dim, quant_config=quant_config, prefix=f"{prefix}.k_proj") + self.v_proj = ReplicatedLinear(dim, attn_dim, quant_config=quant_config, prefix=f"{prefix}.v_proj") + self.out_proj = ReplicatedLinear(attn_dim, dim, quant_config=quant_config, prefix=f"{prefix}.out_proj") if qk_norm == "rms_norm": self.norm_q = RMSNorm(self.head_dim, eps=eps) @@ -158,27 +135,24 @@ class DreamXPropeSelfAttention(nn.Module): if self.out_proj.bias is not None: nn.init.zeros_(self.out_proj.bias) - self.attn = LocalAttention( - num_heads=num_heads, - head_size=self.head_dim, - dropout_rate=0, - softmax_scale=None, - causal=False, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA)) + self.attn = LocalAttention(num_heads=num_heads, + head_size=self.head_dim, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA)) - def forward(self, hidden_states: torch.Tensor, - y_camera: dict[str, torch.Tensor]) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, y_camera: dict[str, torch.Tensor]) -> torch.Tensor: if get_sp_world_size() > 1: # The transformer shards the sequence before the block loop and # this branch uses LocalAttention (no all-to-all): under # sequence parallelism each rank would attend only within its # own shard — silently wrong output. Fail loudly until this # path is ported to DistributedAttention and validated. - raise NotImplementedError( - "DreamXPropeSelfAttention does not support sequence " - "parallelism yet (LocalAttention on a sharded sequence " - "corrupts output). Run with sp_size=1.") + raise NotImplementedError("DreamXPropeSelfAttention does not support sequence " + "parallelism yet (LocalAttention on a sharded sequence " + "corrupts output). Run with sp_size=1.") batch_size, seq_len, _ = hidden_states.shape query, _ = self.q_proj(hidden_states) @@ -186,16 +160,13 @@ class DreamXPropeSelfAttention(nn.Module): value, _ = self.v_proj(hidden_states) if self.qk_norm == "rms_norm": - query = query.view(batch_size, seq_len, self.num_heads, - self.head_dim) + query = query.view(batch_size, seq_len, self.num_heads, self.head_dim) key = key.view(batch_size, seq_len, self.num_heads, self.head_dim) query = self.norm_q(query) key = self.norm_k(key) else: - query = self.norm_q(query).view(batch_size, seq_len, - self.num_heads, self.head_dim) - key = self.norm_k(key).view(batch_size, seq_len, self.num_heads, - self.head_dim) + query = self.norm_q(query).view(batch_size, seq_len, self.num_heads, self.head_dim) + key = self.norm_k(key).view(batch_size, seq_len, self.num_heads, self.head_dim) value = value.view(batch_size, seq_len, self.num_heads, self.head_dim) query = query.transpose(1, 2) @@ -210,10 +181,8 @@ class DreamXPropeSelfAttention(nn.Module): intrinsics=y_camera["K"], ) - out = self.attn(query.transpose(1, 2), key.transpose(1, 2), - value.transpose(1, 2)) - out = _dreamx_apply_tiled_projmat(out.transpose(1, 2), - output_projection).transpose(1, 2) + out = self.attn(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)) + out = _dreamx_apply_tiled_projmat(out.transpose(1, 2), output_projection).transpose(1, 2) out = out.flatten(2) out, _ = self.out_proj(out) return out @@ -238,8 +207,7 @@ class DreamXWorldTransformerBlock(WanTransformerBlock): attn_compress: int = 1, cam_self_attn_layers: tuple[int, ...] | None = None, layer_idx: int | None = None): - super().__init__(dim, ffn_dim, num_heads, qk_norm, cross_attn_norm, - eps, added_kv_proj_dim, + super().__init__(dim, ffn_dim, num_heads, qk_norm, cross_attn_norm, eps, added_kv_proj_dim, supported_attention_backends, quant_config, prefix) self.cam_self_attn = None add_cam_attn = add_control_adapter and cam_method == "prope" @@ -248,14 +216,13 @@ class DreamXWorldTransformerBlock(WanTransformerBlock): if add_cam_attn: if num_heads % attn_compress != 0 or dim % attn_compress != 0: raise ValueError("DreamX attn_compress must divide dim and num_heads") - self.cam_self_attn = DreamXPropeSelfAttention( - dim, - dim // attn_compress, - num_heads // attn_compress, - qk_norm=qk_norm, - eps=eps, - quant_config=quant_config, - prefix=f"{prefix}.cam_self_attn") + self.cam_self_attn = DreamXPropeSelfAttention(dim, + dim // attn_compress, + num_heads // attn_compress, + qk_norm=qk_norm, + eps=eps, + quant_config=quant_config, + prefix=f"{prefix}.cam_self_attn") def forward( self, @@ -272,8 +239,7 @@ class DreamXWorldTransformerBlock(WanTransformerBlock): if temb.dim() == 4: shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - self.scale_shift_table.unsqueeze(0) + temb.float()).chunk( - 6, dim=2) + self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2) shift_msa = shift_msa.squeeze(2) scale_msa = scale_msa.squeeze(2) gate_msa = gate_msa.squeeze(2) @@ -282,12 +248,10 @@ class DreamXWorldTransformerBlock(WanTransformerBlock): c_gate_msa = c_gate_msa.squeeze(2) else: e = self.scale_shift_table + temb.float() - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( - 6, dim=1) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1) assert shift_msa.dtype == torch.float32 - norm_hidden_states = (self.norm1(hidden_states.float()) * - (1 + scale_msa) + shift_msa).to(orig_dtype) + norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -312,22 +276,17 @@ class DreamXWorldTransformerBlock(WanTransformerBlock): attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) if self.cam_self_attn is not None and y_camera is not None: - attn_output = attn_output + self.cam_self_attn( - norm_hidden_states, y_camera) + attn_output = attn_output + self.cam_self_attn(norm_hidden_states, y_camera) null_shift = null_scale = torch.tensor([0], device=hidden_states.device) - norm_hidden_states, hidden_states = self.self_attn_residual_norm( - hidden_states, attn_output, gate_msa, null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, + null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) - attn_output = self.attn2(norm_hidden_states, - context=encoder_hidden_states, - context_lens=None) - norm_hidden_states, hidden_states = self.cross_attn_residual_norm( - hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) + attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None) + norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa, + c_scale_msa) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) @@ -339,14 +298,12 @@ class DreamXWorldTransformerBlock(WanTransformerBlock): class DreamXWorldTransformer3DModel(WanTransformer3DModel): _fsdp_shard_conditions = DreamXWorldConfig()._fsdp_shard_conditions _compile_conditions = DreamXWorldConfig()._compile_conditions - _supported_attention_backends = DreamXWorldConfig( - )._supported_attention_backends + _supported_attention_backends = DreamXWorldConfig()._supported_attention_backends param_names_mapping = DreamXWorldConfig().param_names_mapping reverse_param_names_mapping = DreamXWorldConfig().reverse_param_names_mapping lora_param_names_mapping = DreamXWorldConfig().lora_param_names_mapping - def __init__(self, config: DreamXWorldConfig, hf_config: dict[str, - Any]) -> None: + def __init__(self, config: DreamXWorldConfig, hf_config: dict[str, Any]) -> None: BaseDiT.__init__(self, config=config, hf_config=hf_config) self.quant_config = config.quant_config @@ -359,7 +316,8 @@ class DreamXWorldTransformer3DModel(WanTransformer3DModel): self.patch_size = config.patch_size self.text_len = config.text_len - assert config.num_attention_heads % get_sp_world_size() == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})" + assert config.num_attention_heads % get_sp_world_size( + ) == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})" self.patch_embedding = PatchEmbed(in_chans=config.in_channels, embed_dim=inner_dim, @@ -372,23 +330,21 @@ class DreamXWorldTransformer3DModel(WanTransformer3DModel): image_embed_dim=config.image_dim, ) self.blocks = nn.ModuleList([ - DreamXWorldTransformerBlock( - inner_dim, - config.ffn_dim, - config.num_attention_heads, - config.qk_norm, - config.cross_attn_norm, - config.eps, - config.added_kv_proj_dim, - self._supported_attention_backends, - quant_config=config.quant_config, - prefix=f"{config.prefix}.blocks.{i}", - add_control_adapter=config.add_control_adapter, - cam_method=config.cam_method, - attn_compress=config.attn_compress, - cam_self_attn_layers=config.cam_self_attn_layers, - layer_idx=i) - for i in range(config.num_layers) + DreamXWorldTransformerBlock(inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + self._supported_attention_backends, + quant_config=config.quant_config, + prefix=f"{config.prefix}.blocks.{i}", + add_control_adapter=config.add_control_adapter, + cam_method=config.cam_method, + attn_compress=config.attn_compress, + cam_self_attn_layers=config.cam_self_attn_layers, + layer_idx=i) for i in range(config.num_layers) ]) self.norm_out = LayerNormScaleShift(inner_dim, norm_type="layer", @@ -396,10 +352,8 @@ class DreamXWorldTransformer3DModel(WanTransformer3DModel): elementwise_affine=False, dtype=torch.float32, compute_dtype=torch.float32) - self.proj_out = nn.Linear( - inner_dim, config.out_channels * math.prod(config.patch_size)) - self.scale_shift_table = nn.Parameter( - torch.randn(1, 2, inner_dim) / inner_dim**0.5) + self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) self.gradient_checkpointing = False self.__post_init__() @@ -413,8 +367,7 @@ class DreamXWorldTransformer3DModel(WanTransformer3DModel): y_camera: dict[str, torch.Tensor] | None = None, **kwargs) -> torch.Tensor: orig_dtype = hidden_states.dtype - if encoder_hidden_states is not None and not isinstance( - encoder_hidden_states, torch.Tensor): + if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] if isinstance(encoder_hidden_states_image, list): encoder_hidden_states_image = (encoder_hidden_states_image[0] @@ -428,20 +381,17 @@ class DreamXWorldTransformer3DModel(WanTransformer3DModel): d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] - freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames, post_patch_height, post_patch_width), - self.hidden_size, - self.num_attention_heads, - rope_dim_list, - dtype=torch.float32 if current_platform.is_mps() else torch.float64, - rope_theta=10000) - freqs_cis = (freqs_cos.to(hidden_states.device).float(), - freqs_sin.to(hidden_states.device).float()) + freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000) + freqs_cis = (freqs_cos.to(hidden_states.device).float(), freqs_sin.to(hidden_states.device).float()) hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) - hidden_states, original_seq_len = sequence_model_parallel_shard( - hidden_states, dim=1) + hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1) if timestep.dim() == 2: ts_seq_len = timestep.shape[1] @@ -450,10 +400,7 @@ class DreamXWorldTransformer3DModel(WanTransformer3DModel): ts_seq_len = None temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep, - encoder_hidden_states, - encoder_hidden_states_image, - timestep_seq_len=ts_seq_len) + timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len) if ts_seq_len is not None: timestep_proj = timestep_proj.unflatten(2, (6, -1)) else: @@ -461,8 +408,7 @@ class DreamXWorldTransformer3DModel(WanTransformer3DModel): if encoder_hidden_states_image is not None: if encoder_hidden_states is not None: - encoder_hidden_states = torch.concat( - [encoder_hidden_states_image, encoder_hidden_states], dim=1) + encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) else: encoder_hidden_states = encoder_hidden_states_image @@ -473,33 +419,30 @@ class DreamXWorldTransformer3DModel(WanTransformer3DModel): if torch.is_grad_enabled() and self.gradient_checkpointing: for block in self.blocks: - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, original_seq_len, y_camera) + hidden_states = self._gradient_checkpointing_func(block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, original_seq_len, y_camera) else: for block in self.blocks: - hidden_states = block(hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, original_seq_len, + hidden_states = block(hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, + original_seq_len, y_camera=y_camera) if temb.dim() == 3: - shift, scale = (self.scale_shift_table.unsqueeze(0) + - temb.unsqueeze(2)).chunk(2, dim=2) + shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2) shift = shift.squeeze(2) scale = scale.squeeze(2) else: - shift, scale = (self.scale_shift_table + - temb.unsqueeze(1)).chunk(2, dim=1) + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) hidden_states = self.norm_out(hidden_states, shift, scale) - hidden_states = sequence_model_parallel_all_gather_with_unpad( - hidden_states, original_seq_len, dim=1) + hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, - post_patch_height, - post_patch_width, p_t, p_h, p_w, - -1) + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width, + p_t, p_h, p_w, -1) hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) diff --git a/fastvideo/models/dits/dreamx_world_ar.py b/fastvideo/models/dits/dreamx_world_ar.py index 3ab1db04c..cc001ff40 100644 --- a/fastvideo/models/dits/dreamx_world_ar.py +++ b/fastvideo/models/dits/dreamx_world_ar.py @@ -18,8 +18,7 @@ import torch.nn.functional as F from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig from fastvideo.layers.linear import ReplicatedLinear from fastvideo.models.dits.base import BaseDiT -from fastvideo.models.dits.dreamx_world import (_dreamx_apply_tiled_projmat, - _dreamx_prope_qkv) +from fastvideo.models.dits.dreamx_world import (_dreamx_apply_tiled_projmat, _dreamx_prope_qkv) def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: @@ -40,8 +39,7 @@ def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor return out.transpose(1, 2) -def prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - viewmats: torch.Tensor, Ks: torch.Tensor): +def prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, viewmats: torch.Tensor, Ks: torch.Tensor): q, k, v, output_projection = _dreamx_prope_qkv(q, k, v, viewmats, Ks) def apply_fn_o(x: torch.Tensor) -> torch.Tensor: @@ -54,17 +52,14 @@ def sinusoidal_embedding_1d(dim, position): assert dim % 2 == 0 half = dim // 2 position = position.type(torch.float64) - sinusoid = torch.outer( - position, torch.pow(10000, -torch.arange(half).to(position).div(half))) + sinusoid = torch.outer(position, torch.pow(10000, -torch.arange(half).to(position).div(half))) return torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) def rope_params(max_seq_len, dim, theta=10000): assert dim % 2 == 0 - freqs = torch.outer( - torch.arange(max_seq_len), - 1.0 / torch.pow(theta, - torch.arange(0, dim, 2).to(torch.float64).div(dim))) + freqs = torch.outer(torch.arange(max_seq_len), 1.0 / torch.pow(theta, + torch.arange(0, dim, 2).to(torch.float64).div(dim))) return torch.polar(torch.ones_like(freqs), freqs) @@ -108,6 +103,7 @@ class WanLayerNorm(nn.LayerNorm): class WanCrossAttention(nn.Module): + def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6): super().__init__() self.num_heads = num_heads @@ -168,8 +164,7 @@ def block_relativistic_rope(x, grid_sizes, freqs, start_frame=0, relative_frame_ output = [] for i, (f, h, w) in enumerate(grid_sizes.tolist()): seq_len = f * h * w - x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape( - seq_len, n, -1, 2)) + x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)) if relative_frame_indices is not None: frame_indices = relative_frame_indices.long() @@ -178,10 +173,10 @@ def block_relativistic_rope(x, grid_sizes, freqs, start_frame=0, relative_frame_ freqs_temporal = freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1) freqs_i = torch.cat([ - freqs_temporal, - freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), - freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) - ], dim=-1).reshape(seq_len, 1, -1) + freqs_temporal, freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), freqs[2][:w].view(1, 1, w, -1).expand( + f, h, w, -1) + ], + dim=-1).reshape(seq_len, 1, -1) x_i = torch.view_as_real(x_i * freqs_i).flatten(2) x_i = torch.cat([x_i, x[i, seq_len:]]) @@ -193,8 +188,7 @@ def block_relativistic_rope(x, grid_sizes, freqs, start_frame=0, relative_frame_ class CausalWanSelfAttention(nn.Module): """Self-attention with KV cache and Block-Relativistic RoPE for causal inference.""" - def __init__(self, dim, num_heads, local_attn_size=6, sink_size=1, - qk_norm=True, eps=1e-6): + def __init__(self, dim, num_heads, local_attn_size=6, sink_size=1, qk_norm=True, eps=1e-6): assert dim % num_heads == 0 super().__init__() self.dim = dim @@ -213,8 +207,15 @@ class CausalWanSelfAttention(nn.Module): self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - def forward(self, x, seq_lens, grid_sizes, freqs, kv_cache, - current_start=0, cache_start=None, sink_recache_after_switch=False): + def forward(self, + x, + seq_lens, + grid_sizes, + freqs, + kv_cache, + current_start=0, + cache_start=None, + sink_recache_after_switch=False): """ Args: x: Shape [B, L, C] @@ -269,19 +270,22 @@ class CausalWanSelfAttention(nn.Module): temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len] # Block-Relativistic RoPE: query uses window-relative indices - query_relative_indices = torch.arange( - self.local_attn_size - num_new_frames, self.local_attn_size, device=q.device) - roped_query = block_relativistic_rope( - q, grid_sizes, freqs, relative_frame_indices=query_relative_indices).type_as(v) + query_relative_indices = torch.arange(self.local_attn_size - num_new_frames, + self.local_attn_size, + device=q.device) + roped_query = block_relativistic_rope(q, grid_sizes, freqs, + relative_frame_indices=query_relative_indices).type_as(v) # Block-Relativistic RoPE: cached K uses position-in-window indices num_cache_frames = local_end_index // frame_seqlen cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device) cache_grid_sizes = grid_sizes.clone() cache_grid_sizes[0, 0] = num_cache_frames - roped_temp_k = block_relativistic_rope( - temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, n, d).flatten(1, 2), - cache_grid_sizes, freqs, relative_frame_indices=cache_relative_indices).type_as(v) + roped_temp_k = block_relativistic_rope(temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, + n, d).flatten(1, 2), + cache_grid_sizes, + freqs, + relative_frame_indices=cache_relative_indices).type_as(v) cache_update_info = { "action": "roll_and_insert", @@ -316,18 +320,21 @@ class CausalWanSelfAttention(nn.Module): # RoPE with relative indices (growing sequentially before cache fills) current_frame_in_window = local_start_index // frame_seqlen - query_relative_indices = torch.arange( - current_frame_in_window, current_frame_in_window + num_new_frames, device=q.device) - roped_query = block_relativistic_rope( - q, grid_sizes, freqs, relative_frame_indices=query_relative_indices).type_as(v) + query_relative_indices = torch.arange(current_frame_in_window, + current_frame_in_window + num_new_frames, + device=q.device) + roped_query = block_relativistic_rope(q, grid_sizes, freqs, + relative_frame_indices=query_relative_indices).type_as(v) num_cache_frames = local_end_index // frame_seqlen cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device) cache_grid_sizes = grid_sizes.clone() cache_grid_sizes[0, 0] = num_cache_frames - roped_temp_k = block_relativistic_rope( - temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, n, d).flatten(1, 2), - cache_grid_sizes, freqs, relative_frame_indices=cache_relative_indices).type_as(v) + roped_temp_k = block_relativistic_rope(temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, + n, d).flatten(1, 2), + cache_grid_sizes, + freqs, + relative_frame_indices=cache_relative_indices).type_as(v) cache_update_info = { "action": "direct_insert", @@ -358,10 +365,8 @@ class CausalWanSelfAttention(nn.Module): x = attention(roped_query, k_cat, v_cat) else: window_start = max(0, local_end_index - self.max_attention_size) - x = attention( - roped_query, - roped_temp_k[:, window_start:local_end_index], - temp_v[:, window_start:local_end_index]) + x = attention(roped_query, roped_temp_k[:, window_start:local_end_index], + temp_v[:, window_start:local_end_index]) x = x.flatten(2) x, _ = self.o(x) @@ -371,8 +376,15 @@ class CausalWanSelfAttention(nn.Module): class CausalPropeSelfAttention(nn.Module): """PRoPE self-attention with optional KV cache for camera-controlled inference.""" - def __init__(self, dim, attn_dim, num_heads, window_size=(-1, -1), - local_attn_size=-1, sink_size=0, qk_norm=True, eps=1e-6): + def __init__(self, + dim, + attn_dim, + num_heads, + window_size=(-1, -1), + local_attn_size=-1, + sink_size=0, + qk_norm=True, + eps=1e-6): assert dim % num_heads == 0 assert attn_dim % num_heads == 0 super().__init__() @@ -398,9 +410,18 @@ class CausalPropeSelfAttention(nn.Module): nn.init.zeros_(self.out_proj.weight) nn.init.zeros_(self.out_proj.bias) - def forward(self, x, cam_viewmats, cam_K, seq_lens, grid_sizes, freqs, - kv_cache=None, current_start=0, cache_start=None, - sink_recache_after_switch=False, cache_update_policy="commit_detached"): + def forward(self, + x, + cam_viewmats, + cam_K, + seq_lens, + grid_sizes, + freqs, + kv_cache=None, + current_start=0, + cache_start=None, + sink_recache_after_switch=False, + cache_update_policy="commit_detached"): """ Args: x: Shape [B, L, C] @@ -420,9 +441,11 @@ class CausalPropeSelfAttention(nn.Module): v = v.view(b, s, n, d) # Apply PRoPE (Positional Rotary Position Embedding from camera parameters) - q_t, k_t, v_t, apply_fn_o = prope_qkv( - q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), - viewmats=cam_viewmats, Ks=cam_K) + q_t, k_t, v_t, apply_fn_o = prope_qkv(q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + viewmats=cam_viewmats, + Ks=cam_K) proped_q = q_t.transpose(1, 2) proped_k = k_t.transpose(1, 2) proped_v = v_t.transpose(1, 2) @@ -460,8 +483,12 @@ class CausalPropeSelfAttention(nn.Module): write_len = max(0, local_end_index - write_start_index) if write_len > 0: with torch.no_grad(): - kv_cache["k"][:, write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + write_len].detach() - kv_cache["v"][:, write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + write_len].detach() + kv_cache["k"][:, + write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + + write_len].detach() + kv_cache["v"][:, + write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + + write_len].detach() else: # === DIRECT INSERT MODE === local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() @@ -475,8 +502,12 @@ class CausalPropeSelfAttention(nn.Module): write_len = max(0, local_end_index - write_start_index) if write_len > 0: with torch.no_grad(): - kv_cache["k"][:, write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + write_len].detach() - kv_cache["v"][:, write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + write_len].detach() + kv_cache["k"][:, + write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + + write_len].detach() + kv_cache["v"][:, + write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + + write_len].detach() # Attention: sink tokens + local window if sink_tokens > 0: @@ -495,10 +526,8 @@ class CausalPropeSelfAttention(nn.Module): x_out = attention(proped_q, k_cat, v_cat) else: window_start = max(0, local_end_index - self.max_attention_size) - x_out = attention( - proped_q, - kv_cache["k"][:, window_start:local_end_index].detach(), - kv_cache["v"][:, window_start:local_end_index].detach()) + x_out = attention(proped_q, kv_cache["k"][:, window_start:local_end_index].detach(), + kv_cache["v"][:, window_start:local_end_index].detach()) if not is_recompute and cache_update_policy != "none": kv_cache["global_end_index"].fill_(current_end) @@ -513,8 +542,16 @@ class CausalPropeSelfAttention(nn.Module): class CausalWanAttentionBlock(nn.Module): - def __init__(self, dim, ffn_dim, num_heads, local_attn_size=-1, sink_size=0, - qk_norm=True, cross_attn_norm=False, eps=1e-6, **kwargs): + def __init__(self, + dim, + ffn_dim, + num_heads, + local_attn_size=-1, + sink_size=0, + qk_norm=True, + cross_attn_norm=False, + eps=1e-6, + **kwargs): super().__init__() self.dim = dim self.ffn_dim = ffn_dim @@ -532,10 +569,8 @@ class CausalWanAttentionBlock(nn.Module): # layers self.norm1 = WanLayerNorm(dim, eps) - self.self_attn = CausalWanSelfAttention( - dim, num_heads, local_attn_size, sink_size, qk_norm, eps) - self.norm3 = WanLayerNorm( - dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.self_attn = CausalWanSelfAttention(dim, num_heads, local_attn_size, sink_size, qk_norm, eps) + self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() self.cross_attn = WanCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps) self.norm2 = WanLayerNorm(dim, eps) # nn.Linear (not ReplicatedLinear) on purpose: the official checkpoint @@ -543,36 +578,48 @@ class CausalWanAttentionBlock(nn.Module): # copy-only converter and the strict-load tests require verbatim, and # ReplicatedLinear's (out, bias) tuple return cannot compose inside # nn.Sequential without renaming the state-dict surface. - self.ffn = nn.Sequential( - nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), - nn.Linear(ffn_dim, dim)) + self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), nn.Linear(ffn_dim, dim)) # PRoPE self-attention branch for camera control add_cam_attn = self.add_control_adapter and self.cam_method == 'prope' if add_cam_attn and cam_self_attn_layers is not None: add_cam_attn = self.layer_idx in cam_self_attn_layers if add_cam_attn: - self.cam_self_attn = CausalPropeSelfAttention( - dim, dim // self.attn_compress, num_heads, - local_attn_size=local_attn_size, sink_size=sink_size, - qk_norm=qk_norm, eps=eps) + self.cam_self_attn = CausalPropeSelfAttention(dim, + dim // self.attn_compress, + num_heads, + local_attn_size=local_attn_size, + sink_size=sink_size, + qk_norm=qk_norm, + eps=eps) # modulation self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) - def forward(self, x, e, seq_lens, grid_sizes, freqs, context, context_lens, - kv_cache, crossattn_cache=None, current_start=0, cache_start=None, - cam_viewmats=None, cam_K=None, sink_recache_after_switch=False, + def forward(self, + x, + e, + seq_lens, + grid_sizes, + freqs, + context, + context_lens, + kv_cache, + crossattn_cache=None, + current_start=0, + cache_start=None, + cam_viewmats=None, + cam_K=None, + sink_recache_after_switch=False, cache_update_policy="commit_detached"): num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1] e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2) # self-attention - attn_input = (self.norm1(x).unflatten( - dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]).flatten(1, 2) - y, cache_update_info = self.self_attn( - attn_input, seq_lens, grid_sizes, freqs, kv_cache, - current_start, cache_start, sink_recache_after_switch) + attn_input = (self.norm1(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]).flatten( + 1, 2) + y, cache_update_info = self.self_attn(attn_input, seq_lens, grid_sizes, freqs, kv_cache, current_start, + cache_start, sink_recache_after_switch) # PRoPE camera attention (parallel branch) if hasattr(self, 'cam_self_attn') and cam_viewmats is not None and cam_K is not None: @@ -584,19 +631,23 @@ class CausalWanAttentionBlock(nn.Module): "global_end_index": kv_cache["prope_global_end_index"], "local_end_index": kv_cache["prope_local_end_index"], } - y = y + self.cam_self_attn( - attn_input, cam_viewmats, cam_K, seq_lens, grid_sizes, freqs, - kv_cache=prope_kv_cache, current_start=current_start, - cache_start=cache_start, cache_update_policy=cache_update_policy) + y = y + self.cam_self_attn(attn_input, + cam_viewmats, + cam_K, + seq_lens, + grid_sizes, + freqs, + kv_cache=prope_kv_cache, + current_start=current_start, + cache_start=cache_start, + cache_update_policy=cache_update_policy) x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[2]).flatten(1, 2) # cross-attention & FFN - x = x + self.cross_attn(self.norm3(x), context, context_lens, - crossattn_cache=crossattn_cache) + x = x + self.cross_attn(self.norm3(x), context, context_lens, crossattn_cache=crossattn_cache) y = self.ffn( - (self.norm2(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) - * (1 + e[4]) + e[3]).flatten(1, 2)) + (self.norm2(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[4]) + e[3]).flatten(1, 2)) x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[5]).flatten(1, 2) return x, cache_update_info @@ -619,9 +670,7 @@ class CausalHead(nn.Module): def forward(self, x, e): num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1] e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2) - x, _ = self.head( - self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) - * (1 + e[1]) + e[0]) + x, _ = self.head(self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]) return x @@ -680,27 +729,26 @@ class DreamXWorldARTransformer3DModel(BaseDiT): # embeddings — nn.Linear inside nn.Sequential on purpose: the official # checkpoint keys are positional (text_embedding.0/.2, time_embedding.0/.2, # time_projection.1) and must load verbatim (see ffn comment above). - self.patch_embedding = nn.Conv3d( - in_dim, dim, kernel_size=patch_size, stride=patch_size) - self.text_embedding = nn.Sequential( - nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), - nn.Linear(dim, dim)) - self.time_embedding = nn.Sequential( - nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) - self.time_projection = nn.Sequential( - nn.SiLU(), nn.Linear(dim, dim * 6)) + self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential(nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), nn.Linear(dim, dim)) + self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) # transformer blocks self.blocks = nn.ModuleList([ - CausalWanAttentionBlock( - dim, ffn_dim, num_heads, local_attn_size, sink_size, - qk_norm, cross_attn_norm, eps, - add_control_adapter=add_control_adapter, - cam_method=cam_method, - attn_compress=attn_compress, - layer_idx=layer_idx, - cam_self_attn_layers=cam_self_attn_layers) - for layer_idx in range(num_layers) + CausalWanAttentionBlock(dim, + ffn_dim, + num_heads, + local_attn_size, + sink_size, + qk_norm, + cross_attn_norm, + eps, + add_control_adapter=add_control_adapter, + cam_method=cam_method, + attn_compress=attn_compress, + layer_idx=layer_idx, + cam_self_attn_layers=cam_self_attn_layers) for layer_idx in range(num_layers) ]) for layer_idx, block in enumerate(self.blocks): block.self_attn.layer_idx = layer_idx @@ -712,11 +760,11 @@ class DreamXWorldARTransformer3DModel(BaseDiT): # RoPE frequencies assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 d = dim // num_heads - self.freqs = torch.cat([ - rope_params(1024, d - 4 * (d // 6)), - rope_params(1024, 2 * (d // 6)), - rope_params(1024, 2 * (d // 6)) - ], dim=1) + self.freqs = torch.cat( + [rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6))], + dim=1) self.num_attention_heads = num_heads self.attention_head_dim = dim // num_heads @@ -728,10 +776,22 @@ class DreamXWorldARTransformer3DModel(BaseDiT): self.num_frame_per_block = config.arch_config.num_frames_per_block self.__post_init__() - def forward(self, x=None, t=None, context=None, seq_len=None, y=None, y_camera=None, - kv_cache=None, crossattn_cache=None, current_start=0, - cache_start=0, cache_update_policy="commit_detached", - hidden_states=None, encoder_hidden_states=None, timestep=None, **kwargs): + def forward(self, + x=None, + t=None, + context=None, + seq_len=None, + y=None, + y_camera=None, + kv_cache=None, + crossattn_cache=None, + current_start=0, + cache_start=0, + cache_update_policy="commit_detached", + hidden_states=None, + encoder_hidden_states=None, + timestep=None, + **kwargs): """ Causal inference with KV caching. See Algorithm 2 of CausVid (https://arxiv.org/abs/2412.07772). @@ -766,7 +826,8 @@ class DreamXWorldARTransformer3DModel(BaseDiT): seq_len = int(t.shape[1]) if t.dim() > 1 else int(t.numel()) elif x is not None: sample = x[0] - seq_len = (sample.shape[1] // self.patch_size[0]) * (sample.shape[2] // self.patch_size[1]) * (sample.shape[3] // self.patch_size[2]) + seq_len = (sample.shape[1] // self.patch_size[0]) * (sample.shape[2] // self.patch_size[1]) * ( + sample.shape[3] // self.patch_size[2]) if x is None or t is None or context is None or seq_len is None: raise ValueError("DreamXWorldARTransformer3DModel requires x/t/context/seq_len or FastVideo aliases") @@ -777,34 +838,28 @@ class DreamXWorldARTransformer3DModel(BaseDiT): rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6)), - ], dim=1).to(device) + ], + dim=1).to(device) if y is not None: x = [torch.cat([u, v], dim=0) for u, v in zip(x, y, strict=True)] # patch embedding x = [self.patch_embedding(u.unsqueeze(0)) for u in x] - grid_sizes = torch.stack( - [torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) x = [u.flatten(2).transpose(1, 2) for u in x] seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) assert seq_lens.max() <= seq_len x = torch.cat(x) # time embedding - e = self.time_embedding( - sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x)) - e0 = self.time_projection(e).unflatten( - 1, (6, self.dim)).unflatten(dim=0, sizes=t.shape) + e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x)) + e0 = self.time_projection(e).unflatten(1, (6, self.dim)).unflatten(dim=0, sizes=t.shape) # text embedding context_lens = None context = self.text_embedding( - torch.stack([ - torch.cat( - [u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) - for u in context - ])) + torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context])) # camera parameters if y_camera is not None and isinstance(y_camera, dict): @@ -815,9 +870,14 @@ class DreamXWorldARTransformer3DModel(BaseDiT): cam_K = None block_kwargs = dict( - e=e0, seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=self.freqs, - context=context, context_lens=context_lens, - cam_viewmats=cam_viewmats, cam_K=cam_K, + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=self.freqs, + context=context, + context_lens=context_lens, + cam_viewmats=cam_viewmats, + cam_K=cam_K, cache_update_policy=cache_update_policy, ) @@ -868,7 +928,8 @@ class DreamXWorldARTransformer3DModel(BaseDiT): cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - if write_end_index > write_start_index and new_k.shape[1] == (write_end_index - write_start_index): + if write_end_index > write_start_index and new_k.shape[1] == (write_end_index - + write_start_index): cache["k"][:, write_start_index:write_end_index] = new_k cache["v"][:, write_start_index:write_end_index] = new_v @@ -878,7 +939,8 @@ class DreamXWorldARTransformer3DModel(BaseDiT): new_k = update_info["new_k"].detach() new_v = update_info["new_v"].detach() - if write_end_index > write_start_index and new_k.shape[1] == (write_end_index - write_start_index): + if write_end_index > write_start_index and new_k.shape[1] == (write_end_index - + write_start_index): cache["k"][:, write_start_index:write_end_index] = new_k cache["v"][:, write_start_index:write_end_index] = new_v diff --git a/fastvideo/models/dits/flux.py b/fastvideo/models/dits/flux.py index 0e30ad15d..df33fc3c6 100644 --- a/fastvideo/models/dits/flux.py +++ b/fastvideo/models/dits/flux.py @@ -66,6 +66,7 @@ class FluxPosEmbed(nn.Module): class FluxCombinedTimestepGuidanceTextProjEmbeddings(nn.Module): + def __init__(self, embedding_dim: int, pooled_projection_dim: int) -> None: super().__init__() self.time_proj = Timesteps( @@ -105,6 +106,7 @@ class FluxCombinedTimestepGuidanceTextProjEmbeddings(nn.Module): class FluxAdaLayerNormZeroSingle(nn.Module): + def __init__(self, embedding_dim: int, bias: bool = True) -> None: super().__init__() self.silu = nn.SiLU() @@ -149,12 +151,10 @@ class FluxJointAttention(nn.Module): self.add_k_proj = ReplicatedLinear(dim, self.inner_dim, bias=True) self.add_v_proj = ReplicatedLinear(dim, self.inner_dim, bias=True) - self.to_out = nn.ModuleList( - [ - ReplicatedLinear(self.inner_dim, dim, bias=True), - nn.Dropout(0.0), - ] - ) + self.to_out = nn.ModuleList([ + ReplicatedLinear(self.inner_dim, dim, bias=True), + nn.Dropout(0.0), + ]) self.to_add_out = ReplicatedLinear(self.inner_dim, dim, bias=True) self.attn = DistributedAttention( @@ -259,6 +259,7 @@ class FluxSingleStreamAttention(nn.Module): class FluxTransformerBlock(nn.Module): + def __init__( self, dim: int, @@ -298,9 +299,8 @@ class FluxTransformerBlock(nn.Module): ) -> tuple[torch.Tensor, torch.Tensor]: del joint_attention_kwargs norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) - (norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp) = self.norm1_context( - encoder_hidden_states, emb=temb - ) + (norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, + c_gate_mlp) = self.norm1_context(encoder_hidden_states, emb=temb) attn_output, context_attn_output = self.attn( hidden_states=norm_hidden_states, @@ -330,6 +330,7 @@ class FluxTransformerBlock(nn.Module): class FluxSingleTransformerBlock(nn.Module): + def __init__( self, dim: int, @@ -384,8 +385,8 @@ class FluxTransformer2DModel(BaseDiT): """FastVideo FLUX transformer; load Diffusers FLUX safetensors 1:1.""" _fsdp_shard_conditions = [ - lambda n, m: (n.startswith("transformer_blocks.") or n.startswith("single_transformer_blocks.")) - and n.split(".")[-1].isdigit(), + lambda n, m: (n.startswith("transformer_blocks.") or n.startswith("single_transformer_blocks.")) and n.split( + ".")[-1].isdigit(), ] _compile_conditions = _fsdp_shard_conditions # HF weight names already match this module layout (cf. SGLang regex maps). @@ -426,28 +427,22 @@ class FluxTransformer2DModel(BaseDiT): self.context_embedder = ReplicatedLinear(arch.joint_attention_dim, self.inner_dim) self.x_embedder = ReplicatedLinear(arch.in_channels, self.inner_dim) - self.transformer_blocks = nn.ModuleList( - [ - FluxTransformerBlock( - dim=self.inner_dim, - num_attention_heads=arch.num_attention_heads, - attention_head_dim=arch.attention_head_dim, - supported_attention_backends=self._supported_attention_backends, - ) - for _ in range(arch.num_layers) - ] - ) - self.single_transformer_blocks = nn.ModuleList( - [ - FluxSingleTransformerBlock( - dim=self.inner_dim, - num_attention_heads=arch.num_attention_heads, - attention_head_dim=arch.attention_head_dim, - supported_attention_backends=self._supported_attention_backends, - ) - for _ in range(arch.num_single_layers) - ] - ) + self.transformer_blocks = nn.ModuleList([ + FluxTransformerBlock( + dim=self.inner_dim, + num_attention_heads=arch.num_attention_heads, + attention_head_dim=arch.attention_head_dim, + supported_attention_backends=self._supported_attention_backends, + ) for _ in range(arch.num_layers) + ]) + self.single_transformer_blocks = nn.ModuleList([ + FluxSingleTransformerBlock( + dim=self.inner_dim, + num_attention_heads=arch.num_attention_heads, + attention_head_dim=arch.attention_head_dim, + supported_attention_backends=self._supported_attention_backends, + ) for _ in range(arch.num_single_layers) + ]) self.norm_out = SD3AdaLayerNormContinuous( self.inner_dim, @@ -571,7 +566,7 @@ class FluxTransformer2DModel(BaseDiT): output, _ = self.proj_out(hidden_states) if not return_dict: - return (output,) + return (output, ) return FluxTransformer2DModelOutput(sample=output) diff --git a/fastvideo/models/dits/flux_2.py b/fastvideo/models/dits/flux_2.py index d5ab559c2..de1ea4cb4 100644 --- a/fastvideo/models/dits/flux_2.py +++ b/fastvideo/models/dits/flux_2.py @@ -21,6 +21,7 @@ logger = init_logger(__name__) # pylint: disable=invalid-name class TimestepEmbedding(nn.Module): + def __init__( self, in_channels: int, @@ -51,6 +52,7 @@ class TimestepEmbedding(nn.Module): class AdaLayerNormContinuous(nn.Module): + def __init__( self, embedding_dim: int, @@ -90,9 +92,7 @@ def apply_qk_norm( return q_out, k_out -def _get_qkv_projections( - attn: "Flux2Attention", hidden_states, encoder_hidden_states=None -): +def _get_qkv_projections(attn: "Flux2Attention", hidden_states, encoder_hidden_states=None): query, _ = attn.to_q(hidden_states) key, _ = attn.to_k(hidden_states) value, _ = attn.to_v(hidden_states) @@ -150,12 +150,8 @@ class Flux2FeedForward(nn.Module): self.linear_in = nn.Linear(dim, inner_dim * 2, bias=bias) self.linear_out = nn.Linear(inner_dim, dim_out, bias=bias) else: - self.linear_in = ColumnParallelLinear( - dim, inner_dim * 2, bias=bias, gather_output=True - ) - self.linear_out = ColumnParallelLinear( - inner_dim, dim_out, bias=bias, gather_output=True - ) + self.linear_in = ColumnParallelLinear(dim, inner_dim * 2, bias=bias, gather_output=True) + self.linear_out = ColumnParallelLinear(inner_dim, dim_out, bias=bias, gather_output=True) self.act_fn = Flux2SwiGLU(compute_in_fp32=swiglu_fp32) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -211,39 +207,21 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin): self.added_kv_proj_dim = added_kv_proj_dim self.added_proj_bias = added_proj_bias - self.to_q = ColumnParallelLinear( - query_dim, self.inner_dim, bias=bias, gather_output=True - ) - self.to_k = ColumnParallelLinear( - query_dim, self.inner_dim, bias=bias, gather_output=True - ) - self.to_v = ColumnParallelLinear( - query_dim, self.inner_dim, bias=bias, gather_output=True - ) + self.to_q = ColumnParallelLinear(query_dim, self.inner_dim, bias=bias, gather_output=True) + self.to_k = ColumnParallelLinear(query_dim, self.inner_dim, bias=bias, gather_output=True) + self.to_v = ColumnParallelLinear(query_dim, self.inner_dim, bias=bias, gather_output=True) # QK norm: match Diffusers ``nn.RMSNorm`` (not vLLM-style fp32 variance path). - self.norm_q = nn.RMSNorm( - dim_head, eps=eps, elementwise_affine=elementwise_affine - ) - self.norm_k = nn.RMSNorm( - dim_head, eps=eps, elementwise_affine=elementwise_affine - ) + self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + self.norm_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) self.to_out = torch.nn.ModuleList([]) - self.to_out.append( - ColumnParallelLinear( - self.inner_dim, self.out_dim, bias=out_bias, gather_output=True - ) - ) + self.to_out.append(ColumnParallelLinear(self.inner_dim, self.out_dim, bias=out_bias, gather_output=True)) self.to_out.append(torch.nn.Dropout(dropout)) if added_kv_proj_dim is not None: - self.norm_added_q = nn.RMSNorm( - dim_head, eps=eps, elementwise_affine=elementwise_affine - ) - self.norm_added_k = nn.RMSNorm( - dim_head, eps=eps, elementwise_affine=elementwise_affine - ) + self.norm_added_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + self.norm_added_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) self.add_q_proj = ColumnParallelLinear( added_kv_proj_dim, self.inner_dim, @@ -262,9 +240,7 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin): bias=added_proj_bias, gather_output=True, ) - self.to_add_out = ColumnParallelLinear( - self.inner_dim, query_dim, bias=out_bias, gather_output=True - ) + self.to_add_out = ColumnParallelLinear(self.inner_dim, query_dim, bias=out_bias, gather_output=True) self.attn = LocalAttention( num_heads=num_heads, @@ -280,9 +256,8 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin): encoder_hidden_states: Optional[torch.Tensor] = None, freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: - query, key, value, encoder_query, encoder_key, encoder_value = ( - _get_qkv_projections(self, hidden_states, encoder_hidden_states) - ) + query, key, value, encoder_query, encoder_key, encoder_value = (_get_qkv_projections( + self, hidden_states, encoder_hidden_states)) query = query.unflatten(-1, (self.heads, -1)) key = key.unflatten(-1, (self.heads, -1)) @@ -411,12 +386,8 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): ) self.mlp_act_fn = Flux2SwiGLU(compute_in_fp32=False) - self.norm_q = nn.RMSNorm( - dim_head, eps=eps, elementwise_affine=elementwise_affine - ) - self.norm_k = nn.RMSNorm( - dim_head, eps=eps, elementwise_affine=elementwise_affine - ) + self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + self.norm_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) # Fused attention output projection + MLP output projection self.to_out = ColumnParallelLinear( @@ -464,12 +435,8 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): cos = cos.to(device=query.device) sin = sin.to(device=query.device) # Single stream: match diffusers sequence_dim=1 (query/key [B, S, H, D], no transpose) - query = apply_rotary_emb( - query, (cos, sin), use_real=True, use_real_unbind_dim=-1, sequence_dim=1 - ) - key = apply_rotary_emb( - key, (cos, sin), use_real=True, use_real_unbind_dim=-1, sequence_dim=1 - ) + query = apply_rotary_emb(query, (cos, sin), use_real=True, use_real_unbind_dim=-1, sequence_dim=1) + key = apply_rotary_emb(key, (cos, sin), use_real=True, use_real_unbind_dim=-1, sequence_dim=1) hidden_states = self.attn(query, key, value) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) @@ -485,6 +452,7 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): class Flux2SingleTransformerBlock(nn.Module): + def __init__( self, dim: int, @@ -558,6 +526,7 @@ class Flux2SingleTransformerBlock(nn.Module): class Flux2TransformerBlock(nn.Module): + def __init__( self, dim: int, @@ -604,21 +573,15 @@ class Flux2TransformerBlock(nn.Module): self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - temb_mod_params_img: Tuple[ - Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ... - ], - temb_mod_params_txt: Tuple[ - Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ... - ], + temb_mod_params_img: Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...], + temb_mod_params_txt: Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...], freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: joint_attention_kwargs = joint_attention_kwargs or {} # Modulation parameters shape: [1, 1, self.dim] - (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = ( - temb_mod_params_img - ) + (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = (temb_mod_params_img) (c_shift_msa, c_scale_msa, c_gate_msa), ( c_shift_mlp, c_scale_mlp, @@ -631,9 +594,7 @@ class Flux2TransformerBlock(nn.Module): # Conditioning txt stream norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states) - norm_encoder_hidden_states = ( - 1 + c_scale_msa - ) * norm_encoder_hidden_states + c_shift_msa + norm_encoder_hidden_states = (1 + c_scale_msa) * norm_encoder_hidden_states + c_shift_msa # Attention on concatenated img + txt stream attention_outputs = self.attn( @@ -659,9 +620,7 @@ class Flux2TransformerBlock(nn.Module): encoder_hidden_states = encoder_hidden_states + context_attn_output norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) - norm_encoder_hidden_states = ( - norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp - ) + norm_encoder_hidden_states = (norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp) context_ff_output = self.ff_context(norm_encoder_hidden_states) context_ff_update = c_gate_mlp * context_ff_output @@ -673,6 +632,7 @@ class Flux2TransformerBlock(nn.Module): class Flux2TimestepGuidanceEmbeddings(nn.Module): + def __init__( self, in_channels: int = 256, @@ -682,12 +642,10 @@ class Flux2TimestepGuidanceEmbeddings(nn.Module): ): super().__init__() - self.time_proj = Timesteps( - num_channels=in_channels, flip_sin_to_cos=True, downscale_freq_shift=0 - ) - self.timestep_embedder = TimestepEmbedding( - in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias - ) + self.time_proj = Timesteps(num_channels=in_channels, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=in_channels, + time_embed_dim=embedding_dim, + sample_proj_bias=bias) if guidance_embeds: self.guidance_embedder = TimestepEmbedding( @@ -698,19 +656,13 @@ class Flux2TimestepGuidanceEmbeddings(nn.Module): else: self.guidance_embedder = None - def forward( - self, timestep: torch.Tensor, guidance: Optional[torch.Tensor] = None - ) -> torch.Tensor: + def forward(self, timestep: torch.Tensor, guidance: Optional[torch.Tensor] = None) -> torch.Tensor: timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder( - timesteps_proj.to(timestep.dtype) - ) # (N, D) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(timestep.dtype)) # (N, D) if guidance is not None and self.guidance_embedder is not None: guidance_proj = self.time_proj(guidance) - guidance_emb = self.guidance_embedder( - guidance_proj.to(guidance.dtype) - ) # (N, D) + guidance_emb = self.guidance_embedder(guidance_proj.to(guidance.dtype)) # (N, D) time_guidance_emb = timesteps_emb + guidance_emb return time_guidance_emb else: @@ -718,18 +670,15 @@ class Flux2TimestepGuidanceEmbeddings(nn.Module): class Flux2Modulation(nn.Module): + def __init__(self, dim: int, mod_param_sets: int = 2, bias: bool = False): super().__init__() self.mod_param_sets = mod_param_sets - self.linear = ColumnParallelLinear( - dim, dim * 3 * self.mod_param_sets, bias=bias, gather_output=True - ) + self.linear = ColumnParallelLinear(dim, dim * 3 * self.mod_param_sets, bias=bias, gather_output=True) self.act_fn = nn.SiLU() - def forward( - self, temb: torch.Tensor - ) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]: + def forward(self, temb: torch.Tensor) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]: mod = self.act_fn(temb) mod, _ = self.linear(mod) @@ -737,9 +686,7 @@ class Flux2Modulation(nn.Module): mod = mod.unsqueeze(1) mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1) # Return tuple of 3-tuples of modulation params shift/scale/gate - return tuple( - mod_params[3 * i : 3 * (i + 1)] for i in range(self.mod_param_sets) - ) + return tuple(mod_params[3 * i:3 * (i + 1)] for i in range(self.mod_param_sets)) def _flux2_get_1d_rotary_pos_embed_diffusers( @@ -755,7 +702,7 @@ def _flux2_get_1d_rotary_pos_embed_diffusers( pos = pos.reshape(-1).float() theta_f = float(theta) * ntk_factor inv = torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device) - freqs = (1.0 / (theta_f ** (inv / float(dim)))) / linear_factor + freqs = (1.0 / (theta_f**(inv / float(dim)))) / linear_factor freqs = torch.outer(pos, freqs) freqs_cos = freqs.cos().repeat_interleave(2, dim=-1).float() freqs_sin = freqs.sin().repeat_interleave(2, dim=-1).float() @@ -769,9 +716,7 @@ def _flux2_pad_ids_to_axes(ids: torch.Tensor, n_axes: int) -> torch.Tensor: if ids.shape[-1] > n_axes: return ids[..., :n_axes] need = n_axes - ids.shape[-1] - return torch.cat( - [ids, ids[..., -1:].expand(*ids.shape[:-1], need)], dim=-1 - ) + return torch.cat([ids, ids[..., -1:].expand(*ids.shape[:-1], need)], dim=-1) class Flux2PosEmbed(nn.Module): @@ -791,9 +736,7 @@ class Flux2PosEmbed(nn.Module): pos = ids.float() is_mps = pos.device.type == "mps" is_npu = pos.device.type == "npu" - freqs_dtype = ( - torch.float32 if (is_mps or is_npu) else torch.float64 - ) + freqs_dtype = (torch.float32 if (is_mps or is_npu) else torch.float64) cos_out: List[torch.Tensor] = [] sin_out: List[torch.Tensor] = [] for i in range(len(self.axes_dim)): @@ -890,61 +833,46 @@ class Flux2Transformer2DModel(BaseDiT): # 3. Modulation (double stream and single stream blocks share modulation parameters, resp.) # Two sets of shift/scale/gate modulation parameters for the double stream attn and FF sub-blocks - self.double_stream_modulation_img = Flux2Modulation( - self.inner_dim, mod_param_sets=2, bias=False - ) - self.double_stream_modulation_txt = Flux2Modulation( - self.inner_dim, mod_param_sets=2, bias=False - ) + self.double_stream_modulation_img = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False) + self.double_stream_modulation_txt = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False) # Only one set of modulation parameters as the attn and FF sub-blocks are run in parallel for single stream - self.single_stream_modulation = Flux2Modulation( - self.inner_dim, mod_param_sets=1, bias=False - ) + self.single_stream_modulation = Flux2Modulation(self.inner_dim, mod_param_sets=1, bias=False) # 4. Input projections - self.x_embedder = ColumnParallelLinear( - in_channels, self.inner_dim, bias=False, gather_output=True - ) - self.context_embedder = ColumnParallelLinear( - joint_attention_dim, self.inner_dim, bias=False, gather_output=True - ) + self.x_embedder = ColumnParallelLinear(in_channels, self.inner_dim, bias=False, gather_output=True) + self.context_embedder = ColumnParallelLinear(joint_attention_dim, + self.inner_dim, + bias=False, + gather_output=True) # 5. Double Stream Transformer Blocks supported_attention_backends = self._supported_attention_backends - ff_ctx_swiglu = getattr( - config.arch_config, "ff_context_swiglu_fp32", False - ) - self.transformer_blocks = nn.ModuleList( - [ - Flux2TransformerBlock( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - mlp_ratio=mlp_ratio, - eps=eps, - bias=False, - supported_attention_backends=supported_attention_backends, - ff_context_swiglu_fp32=ff_ctx_swiglu, - ) - for _ in range(num_layers) - ] - ) + ff_ctx_swiglu = getattr(config.arch_config, "ff_context_swiglu_fp32", False) + self.transformer_blocks = nn.ModuleList([ + Flux2TransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_ratio=mlp_ratio, + eps=eps, + bias=False, + supported_attention_backends=supported_attention_backends, + ff_context_swiglu_fp32=ff_ctx_swiglu, + ) for _ in range(num_layers) + ]) # 6. Single Stream Transformer Blocks - self.single_transformer_blocks = nn.ModuleList( - [ - Flux2SingleTransformerBlock( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - mlp_ratio=mlp_ratio, - eps=eps, - bias=False, - supported_attention_backends=supported_attention_backends, - ) - for _ in range(num_single_layers) - ] - ) + self.single_transformer_blocks = nn.ModuleList([ + Flux2SingleTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_ratio=mlp_ratio, + eps=eps, + bias=False, + supported_attention_backends=supported_attention_backends, + ) for _ in range(num_single_layers) + ]) # 7. Output layers self.norm_out = AdaLayerNormContinuous( @@ -1010,9 +938,7 @@ class Flux2Transformer2DModel(BaseDiT): if hidden_states.dim() == 5: input_was_5d = True b, c_in, t, h, w = hidden_states.shape - hidden_states = hidden_states.permute(0, 2, 3, 4, 1).reshape( - b, t * h * w, c_in - ) + hidden_states = hidden_states.permute(0, 2, 3, 4, 1).reshape(b, t * h * w, c_in) # 1. Calculate timestep embedding and modulation parameters (match diffusers: scale to [0, 1000]) timestep = timestep.to(hidden_states.dtype) * 1000 @@ -1035,22 +961,27 @@ class Flux2Transformer2DModel(BaseDiT): if input_was_5d: img_h, img_w = h, w else: - raise ValueError( - "Flux 2 4D input requires explicit txt_ids and img_ids " - "with real image dimensions; square-image fallback removed." - ) + raise ValueError("Flux 2 4D input requires explicit txt_ids and img_ids " + "with real image dimensions; square-image fallback removed.") txt_len = encoder_hidden_states.shape[1] txt_ids = torch.cartesian_prod( - torch.arange(1), torch.arange(1), - torch.arange(1), torch.arange(txt_len), + torch.arange(1), + torch.arange(1), + torch.arange(1), + torch.arange(txt_len), ).to(device=hidden_states.device) img_ids = torch.cartesian_prod( - torch.arange(1), torch.arange(img_h), - torch.arange(img_w), torch.arange(1), + torch.arange(1), + torch.arange(img_h), + torch.arange(img_w), + torch.arange(1), ).to(device=hidden_states.device) freqs_cis = compute_flux2_freqs_cis_from_ids( - self.rotary_emb, txt_ids, img_ids, device=hidden_states.device, + self.rotary_emb, + txt_ids, + img_ids, + device=hidden_states.device, ) # 4. Double Stream Transformer Blocks @@ -1084,9 +1015,7 @@ class Flux2Transformer2DModel(BaseDiT): # Reshape 3D (B, T*H*W, C) back to 5D (B, C, T, H, W) for scheduler if input_was_5d: - output = output.reshape(b, t, h, w, self.out_channels).permute( - 0, 4, 1, 2, 3 - ) + output = output.reshape(b, t, h, w, self.out_channels).permute(0, 4, 1, 2, 3) return output diff --git a/fastvideo/models/dits/gen3c.py b/fastvideo/models/dits/gen3c.py index f02fb9e24..e00cf9795 100644 --- a/fastvideo/models/dits/gen3c.py +++ b/fastvideo/models/dits/gen3c.py @@ -81,15 +81,11 @@ class Gen3CPatchEmbed(nn.Module): p_t, p_h, p_w = self.patch_size # Rearrange: b c (t pt) (h ph) (w pw) -> b t h w (c pt ph pw) - hidden_states = hidden_states.reshape( - batch_size, num_channels, - num_frames // p_t, p_t, - height // p_h, p_h, - width // p_w, p_w - ) + hidden_states = hidden_states.reshape(batch_size, num_channels, num_frames // p_t, p_t, height // p_h, p_h, + width // p_w, p_w) hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7) hidden_states = hidden_states.flatten(4, 7) # Flatten patch dimensions - + # Project to model dimension hidden_states, _ = self.proj(hidden_states) return hidden_states @@ -115,13 +111,9 @@ class Gen3CTimestepEmbedding(nn.Module): self.activation = nn.SiLU() if use_adaln_lora: - self.linear_2 = ReplicatedLinear(out_features, - 3 * out_features, - bias=False) + self.linear_2 = ReplicatedLinear(out_features, 3 * out_features, bias=False) else: - self.linear_2 = ReplicatedLinear(out_features, - out_features, - bias=False) + self.linear_2 = ReplicatedLinear(out_features, out_features, bias=False) def forward(self, sample: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: """ @@ -189,7 +181,7 @@ class Gen3CEmbedding(nn.Module): timestep = timestep.flatten() timesteps_proj = self.time_proj(timestep).to(dtype) # (B, D) embedded_timestep, adaln_lora = self.t_embedder(timesteps_proj) - + return embedded_timestep, adaln_lora @@ -256,14 +248,12 @@ class Gen3CSelfAttention(nn.Module): if supported_attention_backends is None: supported_attention_backends = (AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA) - - self.attn = DistributedAttention( - num_heads=num_heads, - head_size=self.head_dim, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix="self_attn" - ) + + self.attn = DistributedAttention(num_heads=num_heads, + head_size=self.head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix="self_attn") def forward( self, @@ -301,10 +291,8 @@ class Gen3CSelfAttention(nn.Module): query = query.transpose(1, 2) # (B, H, S, D_h) -> (B, S, H, D_h) key = key.transpose(1, 2) value = value.transpose(1, 2) - - attn_output, _ = self.attn( - query, key, value, original_seq_len=original_seq_len - ) + + attn_output, _ = self.attn(query, key, value, original_seq_len=original_seq_len) attn_output = attn_output.flatten(-2, -1) # (B, S, H*D_h) # Output projection @@ -343,7 +331,7 @@ class Gen3CCrossAttention(nn.Module): if supported_attention_backends is None: supported_attention_backends = (AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA) - + self.attn = LocalAttention( num_heads=num_heads, head_size=self.head_dim, @@ -445,18 +433,12 @@ class Gen3CTransformerBlock(nn.Module): ReplicatedLinear(adaln_lora_dim, 3 * hidden_size, bias=False), ]) else: - self.adaln_modulation_self_attn = nn.ModuleList([ - nn.SiLU(), - ReplicatedLinear(hidden_size, 3 * hidden_size, bias=False) - ]) - self.adaln_modulation_cross_attn = nn.ModuleList([ - nn.SiLU(), - ReplicatedLinear(hidden_size, 3 * hidden_size, bias=False) - ]) - self.adaln_modulation_mlp = nn.ModuleList([ - nn.SiLU(), - ReplicatedLinear(hidden_size, 3 * hidden_size, bias=False) - ]) + self.adaln_modulation_self_attn = nn.ModuleList( + [nn.SiLU(), ReplicatedLinear(hidden_size, 3 * hidden_size, bias=False)]) + self.adaln_modulation_cross_attn = nn.ModuleList( + [nn.SiLU(), ReplicatedLinear(hidden_size, 3 * hidden_size, bias=False)]) + self.adaln_modulation_mlp = nn.ModuleList( + [nn.SiLU(), ReplicatedLinear(hidden_size, 3 * hidden_size, bias=False)]) def forward( self, @@ -487,24 +469,17 @@ class Gen3CTransformerBlock(nn.Module): # Compute modulation parameters if self.use_adaln_lora and adaln_lora is not None: shift_self_attn, scale_self_attn, gate_self_attn = ( - _run_linear_stack(self.adaln_modulation_self_attn, affine_emb) - + adaln_lora - ).chunk(3, dim=-1) + _run_linear_stack(self.adaln_modulation_self_attn, affine_emb) + adaln_lora).chunk(3, dim=-1) shift_cross_attn, scale_cross_attn, gate_cross_attn = ( - _run_linear_stack(self.adaln_modulation_cross_attn, affine_emb) - + adaln_lora - ).chunk(3, dim=-1) - shift_mlp, scale_mlp, gate_mlp = ( - _run_linear_stack(self.adaln_modulation_mlp, affine_emb) - + adaln_lora - ).chunk(3, dim=-1) + _run_linear_stack(self.adaln_modulation_cross_attn, affine_emb) + adaln_lora).chunk(3, dim=-1) + shift_mlp, scale_mlp, gate_mlp = (_run_linear_stack(self.adaln_modulation_mlp, affine_emb) + + adaln_lora).chunk(3, dim=-1) else: - shift_self_attn, scale_self_attn, gate_self_attn = _run_linear_stack( - self.adaln_modulation_self_attn, affine_emb).chunk(3, dim=-1) - shift_cross_attn, scale_cross_attn, gate_cross_attn = _run_linear_stack( - self.adaln_modulation_cross_attn, affine_emb).chunk(3, dim=-1) - shift_mlp, scale_mlp, gate_mlp = _run_linear_stack( - self.adaln_modulation_mlp, affine_emb).chunk(3, dim=-1) + shift_self_attn, scale_self_attn, gate_self_attn = _run_linear_stack(self.adaln_modulation_self_attn, + affine_emb).chunk(3, dim=-1) + shift_cross_attn, scale_cross_attn, gate_cross_attn = _run_linear_stack(self.adaln_modulation_cross_attn, + affine_emb).chunk(3, dim=-1) + shift_mlp, scale_mlp, gate_mlp = _run_linear_stack(self.adaln_modulation_mlp, affine_emb).chunk(3, dim=-1) # Reshape modulation parameters for broadcasting: (B, D) -> (B, 1, D) shift_self_attn = shift_self_attn.unsqueeze(1).type_as(hidden_states) @@ -553,13 +528,13 @@ class Gen3CRotaryPosEmbed(nn.Module): """ def __init__( - self, - hidden_size: int, - max_size: tuple[int, int, int] = (128, 240, 240), - patch_size: tuple[int, int, int] = (1, 2, 2), - base_fps: int = 24, - rope_scale: tuple[float, float, float] = (1.0, 1.0, 1.0), - enable_fps_modulation: bool = True, + self, + hidden_size: int, + max_size: tuple[int, int, int] = (128, 240, 240), + patch_size: tuple[int, int, int] = (1, 2, 2), + base_fps: int = 24, + rope_scale: tuple[float, float, float] = (1.0, 1.0, 1.0), + enable_fps_modulation: bool = True, ) -> None: super().__init__() @@ -574,13 +549,11 @@ class Gen3CRotaryPosEmbed(nn.Module): self.dim_t = hidden_size - self.dim_h - self.dim_w # NTK-aware extrapolation factors - self.h_ntk_factor = rope_scale[1] ** (self.dim_h / (self.dim_h - 2)) - self.w_ntk_factor = rope_scale[2] ** (self.dim_w / (self.dim_w - 2)) - self.t_ntk_factor = rope_scale[0] ** (self.dim_t / (self.dim_t - 2)) + self.h_ntk_factor = rope_scale[1]**(self.dim_h / (self.dim_h - 2)) + self.w_ntk_factor = rope_scale[2]**(self.dim_w / (self.dim_w - 2)) + self.t_ntk_factor = rope_scale[0]**(self.dim_t / (self.dim_t - 2)) - def forward( - self, hidden_states: torch.Tensor, fps: int | None = None - ) -> tuple[torch.Tensor, torch.Tensor]: + def forward(self, hidden_states: torch.Tensor, fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]: """ Generate 3D RoPE embeddings. """ @@ -594,13 +567,16 @@ class Gen3CRotaryPosEmbed(nn.Module): seq = torch.arange(max(self.max_size), device=device, dtype=torch.float32) - dim_h_range = torch.arange(0, self.dim_h, 2, device=device, dtype=torch.float32)[: (self.dim_h // 2)] / self.dim_h - dim_w_range = torch.arange(0, self.dim_w, 2, device=device, dtype=torch.float32)[: (self.dim_w // 2)] / self.dim_w - dim_t_range = torch.arange(0, self.dim_t, 2, device=device, dtype=torch.float32)[: (self.dim_t // 2)] / self.dim_t + dim_h_range = torch.arange(0, self.dim_h, 2, device=device, + dtype=torch.float32)[:(self.dim_h // 2)] / self.dim_h + dim_w_range = torch.arange(0, self.dim_w, 2, device=device, + dtype=torch.float32)[:(self.dim_w // 2)] / self.dim_w + dim_t_range = torch.arange(0, self.dim_t, 2, device=device, + dtype=torch.float32)[:(self.dim_t // 2)] / self.dim_t - h_spatial_freqs = 1.0 / (h_theta ** dim_h_range) - w_spatial_freqs = 1.0 / (w_theta ** dim_w_range) - temporal_freqs = 1.0 / (t_theta ** dim_t_range) + h_spatial_freqs = 1.0 / (h_theta**dim_h_range) + w_spatial_freqs = 1.0 / (w_theta**dim_w_range) + temporal_freqs = 1.0 / (t_theta**dim_t_range) # Generate positional embeddings half_emb_h = torch.outer(seq[:H], h_spatial_freqs) @@ -622,7 +598,7 @@ class Gen3CRotaryPosEmbed(nn.Module): cos = torch.cos(freqs) sin = torch.sin(freqs) - + return cos, sin @@ -703,9 +679,7 @@ class Gen3CFinalLayer(nn.Module): # Output projection output_dim = out_channels * patch_size[0] * patch_size[1] * patch_size[2] - self.proj_out = ReplicatedLinear(hidden_size, - output_dim, - bias=False) + self.proj_out = ReplicatedLinear(hidden_size, output_dim, bias=False) def forward( self, @@ -723,7 +697,7 @@ class Gen3CFinalLayer(nn.Module): modulation = _run_linear_stack(self.adaln_modulation, affine_emb) if self.use_adaln_lora and adaln_lora is not None: - modulation = modulation + adaln_lora[..., : 2 * self.hidden_size] + modulation = modulation + adaln_lora[..., :2 * self.hidden_size] shift, scale = modulation.chunk(2, dim=-1) @@ -757,7 +731,7 @@ class Gen3CTransformer3DModel(BaseDiT): - Video conditioning with 3D cache buffers - Augment sigma embedding for conditioning noise augmentation """ - + _fsdp_shard_conditions = Gen3CVideoConfig()._fsdp_shard_conditions _compile_conditions = Gen3CVideoConfig()._compile_conditions param_names_mapping = Gen3CVideoConfig().param_names_mapping @@ -780,7 +754,7 @@ class Gen3CTransformer3DModel(BaseDiT): self.adaln_lora_dim = getattr(config, "adaln_lora_dim", 256) self.extra_pos_embed_type = getattr(config, "extra_pos_embed_type", "learnable") self.affine_emb_norm = getattr(config, "affine_emb_norm", True) - + # GEN3C-specific self.frame_buffer_max = getattr(config, "frame_buffer_max", 2) self.buffer_channels = self.frame_buffer_max * 32 # Each buffer: 16 (frame) + 16 (mask) @@ -791,10 +765,8 @@ class Gen3CTransformer3DModel(BaseDiT): patch_embed_in_channels = config.in_channels + 1 + self.buffer_channels if config.concat_padding_mask: patch_embed_in_channels += 1 - - self.patch_embed = Gen3CPatchEmbed( - patch_embed_in_channels, inner_dim, config.patch_size - ) + + self.patch_embed = Gen3CPatchEmbed(patch_embed_in_channels, inner_dim, config.patch_size) # 2. Positional Embeddings self.rope = Gen3CRotaryPosEmbed( @@ -847,8 +819,7 @@ class Gen3CTransformer3DModel(BaseDiT): use_adaln_lora=self.use_adaln_lora, qk_norm=(config.qk_norm == "rms_norm"), supported_attention_backends=config._supported_attention_backends, - ) - for i in range(config.num_layers) + ) for i in range(config.num_layers) ]) # 7. Final Layer @@ -898,10 +869,13 @@ class Gen3CTransformer3DModel(BaseDiT): hidden_states = torch.cat([hidden_states, condition_video_input_mask], dim=1) else: # Default: zeros mask (no conditioning) - condition_video_input_mask = torch.zeros( - batch_size, 1, num_frames, height, width, - device=hidden_states.device, dtype=hidden_states.dtype - ) + condition_video_input_mask = torch.zeros(batch_size, + 1, + num_frames, + height, + width, + device=hidden_states.device, + dtype=hidden_states.dtype) hidden_states = torch.cat([hidden_states, condition_video_input_mask], dim=1) # 2. Concatenate condition_video_pose (3D cache buffers) @@ -909,19 +883,24 @@ class Gen3CTransformer3DModel(BaseDiT): hidden_states = torch.cat([hidden_states, condition_video_pose], dim=1) else: # Default: zeros for pose buffers - pose_zeros = torch.zeros( - batch_size, self.buffer_channels, num_frames, height, width, - device=hidden_states.device, dtype=hidden_states.dtype - ) + pose_zeros = torch.zeros(batch_size, + self.buffer_channels, + num_frames, + height, + width, + device=hidden_states.device, + dtype=hidden_states.dtype) hidden_states = torch.cat([hidden_states, pose_zeros], dim=1) # 3. Concatenate padding mask if needed if self.concat_padding_mask: if padding_mask is None: - padding_mask = torch.ones( - batch_size, 1, height, width, - device=hidden_states.device, dtype=hidden_states.dtype - ) + padding_mask = torch.ones(batch_size, + 1, + height, + width, + device=hidden_states.device, + dtype=hidden_states.dtype) padding_mask = transforms.functional.resize( padding_mask, list(hidden_states.shape[-2:]), @@ -955,9 +934,7 @@ class Gen3CTransformer3DModel(BaseDiT): sp_world_size = get_sp_world_size() original_seq_len = hidden_states.shape[1] if sp_world_size > 1: - hidden_states, original_seq_len = sequence_model_parallel_shard( - hidden_states, dim=1 - ) + hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1) if extra_pos_emb is not None: extra_pos_emb, _ = sequence_model_parallel_shard(extra_pos_emb, dim=1) rope_cos, _ = sequence_model_parallel_shard(rope_emb[0], dim=0) @@ -1006,14 +983,10 @@ class Gen3CTransformer3DModel(BaseDiT): hidden_states = self.final_layer(hidden_states, affine_emb, adaln_lora) if sp_world_size > 1: - hidden_states = sequence_model_parallel_all_gather_with_unpad( - hidden_states, original_seq_len, dim=1 - ) + hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) # 12. Unpatchify: (B, S, P) -> (B, T', H', W', P) -> (B, C, T, H, W) - hidden_states = hidden_states.unflatten( - 1, (post_patch_num_frames, post_patch_height, post_patch_width) - ) + hidden_states = hidden_states.unflatten(1, (post_patch_num_frames, post_patch_height, post_patch_width)) hidden_states = hidden_states.unflatten(-1, (p_t, p_h, p_w, self.out_channels)) hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) hidden_states = hidden_states.flatten(2, 3).flatten(3, 4).flatten(4, 5) diff --git a/fastvideo/models/dits/glm_image.py b/fastvideo/models/dits/glm_image.py index c06ef773f..6471713d2 100644 --- a/fastvideo/models/dits/glm_image.py +++ b/fastvideo/models/dits/glm_image.py @@ -52,9 +52,7 @@ class GlmImageKVCache: def set_mode(self, mode: Optional[str]): if mode is not None and mode not in ["write", "read", "skip"]: - raise ValueError( - f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'" - ) + raise ValueError(f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'") for cache in self.caches: cache.mode = mode @@ -133,18 +131,10 @@ class GlmImageCombinedTimestepSizeEmbeddings(nn.Module): ): super().__init__() - self.time_proj = Timesteps( - num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0 - ) - self.condition_proj = Timesteps( - num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0 - ) - self.timestep_embedder = GlmImageTimestepEmbedding( - in_channels=timesteps_dim, time_embed_dim=embedding_dim - ) - self.condition_embedder = GlmImageTextProjection( - pooled_projection_dim, embedding_dim, act_fn="silu" - ) + self.time_proj = Timesteps(num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.condition_proj = Timesteps(num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = GlmImageTimestepEmbedding(in_channels=timesteps_dim, time_embed_dim=embedding_dim) + self.condition_embedder = GlmImageTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu") def forward( self, @@ -155,21 +145,13 @@ class GlmImageCombinedTimestepSizeEmbeddings(nn.Module): ) -> torch.Tensor: timesteps_proj = self.time_proj(timestep) - crop_coords_proj = self.condition_proj(crop_coords.flatten()).view( - crop_coords.size(0), -1 - ) - target_size_proj = self.condition_proj(target_size.flatten()).view( - target_size.size(0), -1 - ) + crop_coords_proj = self.condition_proj(crop_coords.flatten()).view(crop_coords.size(0), -1) + target_size_proj = self.condition_proj(target_size.flatten()).view(target_size.size(0), -1) condition_proj = torch.cat([crop_coords_proj, target_size_proj], dim=1) - timesteps_emb = self.timestep_embedder( - timesteps_proj.to(dtype=hidden_dtype) - ) - condition_emb = self.condition_embedder( - condition_proj.to(dtype=hidden_dtype) - ) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) + condition_emb = self.condition_embedder(condition_proj.to(dtype=hidden_dtype)) conditioning = timesteps_emb + condition_emb return conditioning @@ -203,9 +185,7 @@ class GlmImageImageProjector(nn.Module): post_patch_width, self.patch_size, ) - hidden_states = ( - hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) - ) + hidden_states = (hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)) hidden_states = self.proj(hidden_states) return hidden_states @@ -231,9 +211,7 @@ class GlmImageAdaLayerNormZero(nn.Module): ) -> Tuple[torch.Tensor, ...]: dtype = hidden_states.dtype norm_hidden_states = self.norm(hidden_states).to(dtype=dtype) - norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to( - dtype=dtype - ) + norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype) emb, _ = self.linear(temb) ( @@ -251,12 +229,8 @@ class GlmImageAdaLayerNormZero(nn.Module): c_gate_mlp, ) = emb.chunk(12, dim=1) - hidden_states = norm_hidden_states * ( - 1 + scale_msa.unsqueeze(1) - ) + shift_msa.unsqueeze(1) - encoder_hidden_states = norm_encoder_hidden_states * ( - 1 + c_scale_msa.unsqueeze(1) - ) + c_shift_msa.unsqueeze(1) + hidden_states = norm_hidden_states * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) + encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_msa.unsqueeze(1)) + c_shift_msa.unsqueeze(1) return ( hidden_states, @@ -303,20 +277,14 @@ class GlmImageAttention(nn.Module): self.to_k = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias) self.to_v = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias) - self.to_out = nn.ModuleList( - [ReplicatedLinear(self.inner_dim, self.out_dim, bias=True)] - ) + self.to_out = nn.ModuleList([ReplicatedLinear(self.inner_dim, self.out_dim, bias=True)]) if qk_norm is None: self.norm_q = None self.norm_k = None elif qk_norm == "layer_norm": - self.norm_q = nn.LayerNorm( - dim_head, eps=eps, elementwise_affine=elementwise_affine - ) - self.norm_k = nn.LayerNorm( - dim_head, eps=eps, elementwise_affine=elementwise_affine - ) + self.norm_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + self.norm_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) else: raise ValueError(f"unknown qk_norm: {qk_norm}") @@ -361,12 +329,14 @@ class GlmImageAttention(nn.Module): if image_rotary_emb is not None: cos, sin = image_rotary_emb - query[:, text_seq_length:, :, :] = _apply_rotary_emb( - query[:, text_seq_length:, :, :], cos, sin, is_neox_style=True - ) - key[:, text_seq_length:, :, :] = _apply_rotary_emb( - key[:, text_seq_length:, :, :], cos, sin, is_neox_style=True - ) + query[:, text_seq_length:, :, :] = _apply_rotary_emb(query[:, text_seq_length:, :, :], + cos, + sin, + is_neox_style=True) + key[:, text_seq_length:, :, :] = _apply_rotary_emb(key[:, text_seq_length:, :, :], + cos, + sin, + is_neox_style=True) # 4. KV Cache handling if kv_cache is not None: @@ -376,9 +346,7 @@ class GlmImageAttention(nn.Module): # Prepend cached condition k/v along seq (dim=1). k_cache, v_cache = kv_cache.get() key = torch.cat([k_cache, key], dim=1) if k_cache is not None else key - value = ( - torch.cat([v_cache, value], dim=1) if v_cache is not None else value - ) + value = (torch.cat([v_cache, value], dim=1) if v_cache is not None else value) elif kv_cache.mode == "skip": pass @@ -390,8 +358,7 @@ class GlmImageAttention(nn.Module): hidden_states, _ = self.to_out[0](hidden_states) encoder_hidden_states, hidden_states = hidden_states.split( - [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1 - ) + [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1) return hidden_states, encoder_hidden_states @@ -428,12 +395,11 @@ class GlmImageTransformerBlock(nn.Module): ) # 2. Feedforward with fused ScaleResidualLayerNorm - self.norm2 = ScaleResidualLayerNormScaleShift( - dim, norm_type="layer", eps=1e-5, elementwise_affine=False - ) - self.norm2_context = ScaleResidualLayerNormScaleShift( - dim, norm_type="layer", eps=1e-5, elementwise_affine=False - ) + self.norm2 = ScaleResidualLayerNormScaleShift(dim, norm_type="layer", eps=1e-5, elementwise_affine=False) + self.norm2_context = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False) self.ff = MLP(input_dim=dim, mlp_hidden_dim=dim * 4, output_dim=dim, act_type="gelu_pytorch_tanh") def forward( @@ -441,12 +407,10 @@ class GlmImageTransformerBlock(nn.Module): hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[ - Union[ - Tuple[torch.Tensor, torch.Tensor], - List[Tuple[torch.Tensor, torch.Tensor]], - ] - ] = None, + image_rotary_emb: Optional[Union[ + Tuple[torch.Tensor, torch.Tensor], + List[Tuple[torch.Tensor, torch.Tensor]], + ]] = None, attention_kwargs: Optional[Dict[str, Any]] = None, kv_cache: Optional[GlmImageLayerKVCache] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -495,9 +459,7 @@ class GlmImageTransformerBlock(nn.Module): ff_output = self.ff(norm_hidden_states) ff_output_context = self.ff(norm_encoder_hidden_states) hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1) - encoder_hidden_states = ( - encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1) - ) + encoder_hidden_states = (encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1)) return hidden_states, encoder_hidden_states @@ -521,30 +483,15 @@ class GlmImageRotaryPosEmbed(nn.Module): width = raw_w // self.patch_size device = hidden_states.device - cache_key = (height, width, device.type, - device.index if device.index is not None else -1) + cache_key = (height, width, device.type, device.index if device.index is not None else -1) if self._cache_key == cache_key and self._cache_value is not None: return self._cache_value dim_h, dim_w = self.dim // 2, self.dim // 2 - h_inv_freq = 1.0 / ( - self.theta - ** ( - torch.arange(0, dim_h, 2, dtype=torch.float32, device=device)[ - : (dim_h // 2) - ] - / dim_h - ) - ) - w_inv_freq = 1.0 / ( - self.theta - ** ( - torch.arange(0, dim_w, 2, dtype=torch.float32, device=device)[ - : (dim_w // 2) - ] - / dim_w - ) - ) + h_inv_freq = 1.0 / (self.theta + **(torch.arange(0, dim_h, 2, dtype=torch.float32, device=device)[:(dim_h // 2)] / dim_h)) + w_inv_freq = 1.0 / (self.theta + **(torch.arange(0, dim_w, 2, dtype=torch.float32, device=device)[:(dim_w // 2)] / dim_w)) h_seq = torch.arange(height, device=device) w_seq = torch.arange(width, device=device) freqs_h = torch.outer(h_seq, h_inv_freq).unsqueeze(1).expand(height, width, -1) @@ -572,9 +519,7 @@ class GlmImageAdaLayerNormContinuous(nn.Module): norm_type: str = "layer_norm", ): super().__init__() - self.linear = nn.Linear( - conditioning_embedding_dim, embedding_dim * 2, bias=bias - ) + self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias) if norm_type == "layer_norm": self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) elif norm_type == "rms_norm": @@ -582,9 +527,7 @@ class GlmImageAdaLayerNormContinuous(nn.Module): else: raise ValueError(f"unknown norm_type {norm_type}") - def forward( - self, x: torch.Tensor, conditioning_embedding: torch.Tensor - ) -> torch.Tensor: + def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: emb = self.linear(conditioning_embedding.to(x.dtype)) scale, shift = torch.chunk(emb, 2, dim=1) x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] @@ -634,23 +577,17 @@ class GlmImageTransformer2DModel(BaseDiT): self.num_channels_latents = arch_config.out_channels # 1. RoPE - self.rotary_emb = GlmImageRotaryPosEmbed( - arch_config.attention_head_dim, arch_config.patch_size, theta=10000.0 - ) + self.rotary_emb = GlmImageRotaryPosEmbed(arch_config.attention_head_dim, arch_config.patch_size, theta=10000.0) # 2. Patch & Text-timestep embedding - self.image_projector = GlmImageImageProjector( - arch_config.in_channels, inner_dim, arch_config.patch_size - ) + self.image_projector = GlmImageImageProjector(arch_config.in_channels, inner_dim, arch_config.patch_size) self.glyph_projector = MLP( input_dim=arch_config.text_embed_dim, mlp_hidden_dim=inner_dim, output_dim=inner_dim, act_type="gelu", ) - self.prior_token_embedding = nn.Embedding( - arch_config.prior_vq_quantizer_codebook_size, inner_dim - ) + self.prior_token_embedding = nn.Embedding(arch_config.prior_vq_quantizer_codebook_size, inner_dim) self.prior_projector = MLP( input_dim=inner_dim, mlp_hidden_dim=inner_dim, @@ -666,24 +603,19 @@ class GlmImageTransformer2DModel(BaseDiT): ) # 3. Transformer blocks - self.transformer_blocks = nn.ModuleList( - [ - GlmImageTransformerBlock( - inner_dim, - arch_config.num_attention_heads, - arch_config.attention_head_dim, - arch_config.time_embed_dim, - supported_attention_backends=self._supported_attention_backends, - prefix=f"transformer_blocks.{i}", - ) - for i in range(arch_config.num_layers) - ] - ) + self.transformer_blocks = nn.ModuleList([ + GlmImageTransformerBlock( + inner_dim, + arch_config.num_attention_heads, + arch_config.attention_head_dim, + arch_config.time_embed_dim, + supported_attention_backends=self._supported_attention_backends, + prefix=f"transformer_blocks.{i}", + ) for i in range(arch_config.num_layers) + ]) # 4. Output projection - self.norm_out = GlmImageAdaLayerNormContinuous( - inner_dim, arch_config.time_embed_dim, elementwise_affine=False - ) + self.norm_out = GlmImageAdaLayerNormContinuous(inner_dim, arch_config.time_embed_dim, elementwise_affine=False) self.proj_out = nn.Linear( inner_dim, arch_config.patch_size * arch_config.patch_size * arch_config.out_channels, @@ -706,12 +638,10 @@ class GlmImageTransformer2DModel(BaseDiT): attention_kwargs: Optional[Dict[str, Any]] = None, kv_caches: Optional[GlmImageKVCache] = None, kv_caches_mode: Optional[str] = None, - freqs_cis: Optional[ - Union[ - Tuple[torch.Tensor, torch.Tensor], - List[Tuple[torch.Tensor, torch.Tensor]], - ] - ] = None, + freqs_cis: Optional[Union[ + Tuple[torch.Tensor, torch.Tensor], + List[Tuple[torch.Tensor, torch.Tensor]], + ]] = None, guidance: torch.Tensor = None, **kwargs, ) -> torch.Tensor: @@ -744,9 +674,7 @@ class GlmImageTransformer2DModel(BaseDiT): prior_hidden_states = self.prior_projector(prior_embedding) hidden_states = hidden_states + prior_hidden_states - temb = self.time_condition_embed( - timestep, target_size, crop_coords, hidden_states.dtype - ) + temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype) temb = F.silu(temb) # 3. Transformer blocks @@ -765,9 +693,7 @@ class GlmImageTransformer2DModel(BaseDiT): hidden_states = self.proj_out(hidden_states) # 5. Unpatchify - hidden_states = hidden_states.reshape( - batch_size, post_patch_height, post_patch_width, -1, p, p - ) + hidden_states = hidden_states.reshape(batch_size, post_patch_height, post_patch_width, -1, p, p) output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) return output.float() diff --git a/fastvideo/models/dits/hunyuangamecraft.py b/fastvideo/models/dits/hunyuangamecraft.py index f993a55e6..ec2def561 100644 --- a/fastvideo/models/dits/hunyuangamecraft.py +++ b/fastvideo/models/dits/hunyuangamecraft.py @@ -36,18 +36,10 @@ class GameCraftFinalLayer(nn.Module): This differs from the HunyuanVideo FinalLayer. """ - def __init__(self, - hidden_size, - patch_size, - out_channels, - dtype=None, - prefix: str = "") -> None: + def __init__(self, hidden_size, patch_size, out_channels, dtype=None, prefix: str = "") -> None: super().__init__() - self.norm_final = nn.LayerNorm(hidden_size, - eps=1e-6, - elementwise_affine=False, - dtype=dtype) + self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False, dtype=dtype) output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels @@ -57,12 +49,11 @@ class GameCraftFinalLayer(nn.Module): params_dtype=dtype, prefix=f"{prefix}.linear") - self.adaLN_modulation = ModulateProjection( - hidden_size, - factor=2, - act_layer="silu", - dtype=dtype, - prefix=f"{prefix}.adaLN_modulation") + self.adaLN_modulation = ModulateProjection(hidden_size, + factor=2, + act_layer="silu", + dtype=dtype, + prefix=f"{prefix}.adaLN_modulation") def forward(self, x, c): # GameCraft uses shift, scale order (verified against official implementation) @@ -78,7 +69,7 @@ class CameraNet(nn.Module): Processes camera parameters (Plücker coordinates) into feature embeddings. """ - + def __init__( self, in_channels: int = 6, @@ -91,44 +82,44 @@ class CameraNet(nn.Module): ): super().__init__() _ = prefix # Unused - - start_channels = in_channels * (downscale_coef ** 2) + + start_channels = in_channels * (downscale_coef**2) input_channels = [start_channels, start_channels // 2, start_channels // 4] self.input_channels = input_channels - + self.unshuffle = nn.PixelUnshuffle(downscale_coef) - + self.encode_first = nn.Sequential( nn.Conv2d(input_channels[0], input_channels[1], kernel_size=1, stride=1, padding=0), nn.GroupNorm(2, input_channels[1]), nn.ReLU(), ) self._initialize_weights(self.encode_first) - + self.encode_second = nn.Sequential( nn.Conv2d(input_channels[1], input_channels[2], kernel_size=1, stride=1, padding=0), nn.GroupNorm(2, input_channels[2]), nn.ReLU(), ) self._initialize_weights(self.encode_second) - + self.final_proj = nn.Conv2d(input_channels[2], out_channels, kernel_size=1) self._zeros_init_linear(self.final_proj) - + self.scale = nn.Parameter(torch.ones(1)) - + self.camera_in = PatchEmbed( patch_size=patch_size, in_chans=out_channels, embed_dim=hidden_size, ) - + def _zeros_init_linear(self, linear): if hasattr(linear, "weight"): nn.init.zeros_(linear.weight) if hasattr(linear, "bias") and linear.bias is not None: nn.init.zeros_(linear.bias) - + def _initialize_weights(self, block): for m in block: if isinstance(m, nn.Conv2d): @@ -136,12 +127,12 @@ class CameraNet(nn.Module): nn.init.normal_(m.weight, mean=0.0, std=np.sqrt(2.0 / n)) if m.bias is not None: nn.init.zeros_(m.bias) - + def compress_time(self, x: torch.Tensor, num_frames: int) -> torch.Tensor: x = rearrange(x, '(b f) c h w -> b f c h w', f=num_frames) batch_size, frames, channels, height, width = x.shape x = rearrange(x, 'b f c h w -> (b h w) c f') - + if x.shape[-1] == 66 or x.shape[-1] == 34: x_len = x.shape[-1] x_clip1 = x[..., :x_len // 2] @@ -159,10 +150,10 @@ class CameraNet(nn.Module): x = torch.cat([x_first[..., None], x_rest], dim=-1) else: x = F.avg_pool1d(x, kernel_size=2, stride=2) - + x = rearrange(x, '(b h w) c f -> (b f) c h w', b=batch_size, h=height, w=width) return x - + def forward(self, camera_states: torch.Tensor) -> torch.Tensor: batch_size, num_frames, channels, height, width = camera_states.shape camera_states = rearrange(camera_states, 'b f c h w -> (b f) c h w') @@ -182,23 +173,23 @@ class HunyuanGameCraftTransformer3DModel(BaseDiT): """ HunyuanGameCraft Transformer - ported from official implementation. """ - + _fsdp_shard_conditions = HunyuanGameCraftArchConfig()._fsdp_shard_conditions _compile_conditions = HunyuanGameCraftArchConfig()._compile_conditions _supported_attention_backends = HunyuanGameCraftArchConfig()._supported_attention_backends param_names_mapping = HunyuanGameCraftConfig().param_names_mapping reverse_param_names_mapping = HunyuanGameCraftConfig().reverse_param_names_mapping - + def __init__(self, config: HunyuanGameCraftConfig, hf_config: dict[str, Any]): super().__init__(config=config, hf_config=hf_config) - + arch = config.arch_config - + if isinstance(arch.patch_size, (list, tuple)): self.patch_size = list(arch.patch_size) else: self.patch_size = [arch.patch_size_t, arch.patch_size, arch.patch_size] - + self.in_channels = arch.in_channels self.out_channels = arch.out_channels self.unpatchify_channels = self.out_channels @@ -212,20 +203,20 @@ class HunyuanGameCraftTransformer3DModel(BaseDiT): self.text_states_dim = arch.text_embed_dim self.text_states_dim_2 = arch.pooled_projection_dim self.dtype = arch.dtype - + pe_dim = self.hidden_size // self.num_heads if sum(self.rope_dim_list) != pe_dim: raise ValueError(f"rope_axes_dim sum {sum(self.rope_dim_list)} != {pe_dim}") - + factory_kwargs = {'dtype': self.dtype} - + self.img_in = PatchEmbed( patch_size=self.patch_size, in_chans=self.in_channels, embed_dim=self.hidden_size, **factory_kwargs, ) - + self.txt_in = SingleTokenRefiner( self.text_states_dim, self.hidden_size, @@ -233,9 +224,9 @@ class HunyuanGameCraftTransformer3DModel(BaseDiT): depth=arch.num_refiner_layers, **factory_kwargs, ) - + self.time_in = TimestepEmbedder(self.hidden_size, **factory_kwargs) - + self.vector_in = MLP( self.text_states_dim_2, self.hidden_size, @@ -243,12 +234,9 @@ class HunyuanGameCraftTransformer3DModel(BaseDiT): act_type="silu", **factory_kwargs, ) - - self.guidance_in = ( - TimestepEmbedder(self.hidden_size, **factory_kwargs) - if self.guidance_embeds else None - ) - + + self.guidance_in = (TimestepEmbedder(self.hidden_size, **factory_kwargs) if self.guidance_embeds else None) + self.double_blocks = nn.ModuleList([ MMDoubleStreamBlock( hidden_size=self.hidden_size, @@ -256,10 +244,9 @@ class HunyuanGameCraftTransformer3DModel(BaseDiT): mlp_ratio=arch.mlp_ratio, supported_attention_backends=self._supported_attention_backends, **factory_kwargs, - ) - for _ in range(arch.num_layers) + ) for _ in range(arch.num_layers) ]) - + self.single_blocks = nn.ModuleList([ MMSingleStreamBlock( hidden_size=self.hidden_size, @@ -267,17 +254,16 @@ class HunyuanGameCraftTransformer3DModel(BaseDiT): mlp_ratio=arch.mlp_ratio, supported_attention_backends=self._supported_attention_backends, **factory_kwargs, - ) - for _ in range(arch.num_single_layers) + ) for _ in range(arch.num_single_layers) ]) - + self.final_layer = GameCraftFinalLayer( self.hidden_size, self.patch_size, self.out_channels, **factory_kwargs, ) - + self.camera_net = CameraNet( in_channels=arch.camera_in_channels, out_channels=16, @@ -285,7 +271,7 @@ class HunyuanGameCraftTransformer3DModel(BaseDiT): patch_size=self.patch_size, hidden_size=self.hidden_size, ) - + def forward( self, x: torch.Tensor, @@ -301,42 +287,41 @@ class HunyuanGameCraftTransformer3DModel(BaseDiT): tt = ot // self.patch_size[0] th = oh // self.patch_size[1] tw = ow // self.patch_size[2] - + text_states = encoder_hidden_states[0] text_states_2 = encoder_hidden_states[1] if len(encoder_hidden_states) > 1 else None text_mask = encoder_attention_mask[0] if encoder_attention_mask else None - + vec = self.time_in(timestep) - + if text_states_2 is not None: vec = vec + self.vector_in(text_states_2) - + if self.guidance_in is not None and guidance is not None: vec = vec + self.guidance_in(guidance) - + img = self.img_in(img) - + if camera_states is not None: latent_len = ot if latent_len == 18: - camera_latents = torch.cat([ - self.camera_net(torch.zeros_like(camera_states)), - self.camera_net(camera_states) - ], dim=1) + camera_latents = torch.cat( + [self.camera_net(torch.zeros_like(camera_states)), + self.camera_net(camera_states)], dim=1) elif latent_len == 9: camera_latents = self.camera_net(camera_states) elif latent_len == 10: - camera_latents = torch.cat([ - self.camera_net(torch.zeros_like(camera_states[:, 0:4, :, :, :])), - self.camera_net(camera_states) - ], dim=1) + camera_latents = torch.cat( + [self.camera_net(torch.zeros_like(camera_states[:, 0:4, :, :, :])), + self.camera_net(camera_states)], + dim=1) else: camera_latents = self.camera_net(camera_states) img = img + camera_latents - + txt = self.txt_in(text_states, timestep) txt_seq_len = txt.shape[1] - + freqs_cos, freqs_sin = get_rotary_pos_embed( (tt, th, tw), self.hidden_size, @@ -347,17 +332,17 @@ class HunyuanGameCraftTransformer3DModel(BaseDiT): freqs_cos = freqs_cos.to(device=img.device, dtype=img.dtype) freqs_sin = freqs_sin.to(device=img.device, dtype=img.dtype) freqs_cis = (freqs_cos, freqs_sin) - + for block in self.double_blocks: img, txt = block(img, txt, vec, freqs_cis) - + x = torch.cat([img, txt], dim=1) - + for block in self.single_blocks: x = block(x, vec, txt_seq_len, freqs_cis) - + img = x[:, :-txt_seq_len, ...] img = self.final_layer(img, vec) img = unpatchify(img, tt, th, tw, self.patch_size, self.out_channels) - + return img diff --git a/fastvideo/models/dits/hunyuanvideo.py b/fastvideo/models/dits/hunyuanvideo.py index 1d25c1ad4..6856d0a5f 100644 --- a/fastvideo/models/dits/hunyuanvideo.py +++ b/fastvideo/models/dits/hunyuanvideo.py @@ -7,14 +7,12 @@ import torch.nn as nn from fastvideo.attention import DistributedAttention, LocalAttention from fastvideo.configs.models.dits import HunyuanVideoConfig -from fastvideo.layers.layernorm import (LayerNormScaleShift, ScaleResidual, - ScaleResidualLayerNormScaleShift) +from fastvideo.layers.layernorm import (LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift) from fastvideo.layers.linear import ReplicatedLinear # TODO(will-PY-refactor): RMSNorm .... from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import get_rotary_pos_embed -from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, - TimestepEmbedder, unpatchify) +from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, TimestepEmbedder, unpatchify) from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum from fastvideo.distributed.communication_op import sequence_model_parallel_shard, sequence_model_parallel_all_gather_with_unpad @@ -111,15 +109,11 @@ class MMDoubleStreamBlock(nn.Module): ) # Fused operations for image stream - self.img_attn_norm = LayerNormScaleShift(hidden_size, - norm_type="layer", - elementwise_affine=False, - dtype=dtype) - self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift( - hidden_size, - norm_type="layer", - elementwise_affine=False, - dtype=dtype) + self.img_attn_norm = LayerNormScaleShift(hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype) + self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(hidden_size, + norm_type="layer", + elementwise_affine=False, + dtype=dtype) self.img_mlp_residual = ScaleResidual() # Image attention components @@ -138,11 +132,7 @@ class MMDoubleStreamBlock(nn.Module): params_dtype=dtype, prefix=f"{prefix}.img_attn_proj") - self.img_mlp = MLP(hidden_size, - mlp_hidden_dim, - bias=True, - dtype=dtype, - prefix=f"{prefix}.img_mlp") + self.img_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, prefix=f"{prefix}.img_mlp") # Text modulation components self.txt_mod = ModulateProjection( @@ -154,41 +144,30 @@ class MMDoubleStreamBlock(nn.Module): ) # Fused operations for text stream - self.txt_attn_norm = LayerNormScaleShift(hidden_size, - norm_type="layer", - elementwise_affine=False, - dtype=dtype) - self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift( - hidden_size, - norm_type="layer", - elementwise_affine=False, - dtype=dtype) + self.txt_attn_norm = LayerNormScaleShift(hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype) + self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(hidden_size, + norm_type="layer", + elementwise_affine=False, + dtype=dtype) self.txt_mlp_residual = ScaleResidual() # Text attention components - self.txt_attn_qkv = ReplicatedLinear(hidden_size, - hidden_size * 3, - bias=True, - params_dtype=dtype) + self.txt_attn_qkv = ReplicatedLinear(hidden_size, hidden_size * 3, bias=True, params_dtype=dtype) # QK norm layers for text self.txt_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) self.txt_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) - self.txt_attn_proj = ReplicatedLinear(hidden_size, - hidden_size, - bias=True, - params_dtype=dtype) + self.txt_attn_proj = ReplicatedLinear(hidden_size, hidden_size, bias=True, params_dtype=dtype) self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype) # Distributed attention - self.attn = DistributedAttention( - num_heads=num_attention_heads, - head_size=head_dim, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.attn") + self.attn = DistributedAttention(num_heads=num_attention_heads, + head_size=head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn") def forward( self, @@ -226,10 +205,8 @@ class MMDoubleStreamBlock(nn.Module): batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1] # Split QKV - img_qkv = img_qkv.view(batch_size, image_seq_len, 3, - self.num_attention_heads, -1) - img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, - 2] + img_qkv = img_qkv.view(batch_size, image_seq_len, 3, self.num_attention_heads, -1) + img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2] # Apply QK-Norm if needed @@ -244,10 +221,8 @@ class MMDoubleStreamBlock(nn.Module): batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1] # Split QKV - txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, - self.num_attention_heads, -1) - txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, - 2] + txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, self.num_attention_heads, -1) + txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2] # Apply QK-Norm if needed txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype) @@ -255,23 +230,21 @@ class MMDoubleStreamBlock(nn.Module): # Run distributed attention img_attn, txt_attn = self.attn(img_q, img_k, img_v, original_seq_len, txt_q, txt_k, txt_v, freqs_cis=freqs_cis) - img_attn_out, _ = self.img_attn_proj( - img_attn.view(batch_size, image_seq_len, -1)) + img_attn_out, _ = self.img_attn_proj(img_attn.view(batch_size, image_seq_len, -1)) # Use fused operation for residual connection, normalization, and modulation - img_mlp_input, img_residual = self.img_attn_residual_mlp_norm( - img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale) + img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(img, img_attn_out, img_attn_gate, img_mlp_shift, + img_mlp_scale) # Process image MLP img_mlp_out = self.img_mlp(img_mlp_input) img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate) # Process text attention output - txt_attn_out, _ = self.txt_attn_proj( - txt_attn.reshape(batch_size, text_seq_len, -1)) + txt_attn_out, _ = self.txt_attn_proj(txt_attn.reshape(batch_size, text_seq_len, -1)) # Use fused operation for residual connection, normalization, and modulation - txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm( - txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale) + txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, + txt_mlp_scale) # Process text MLP txt_mlp_out = self.txt_mlp(txt_mlp_input) @@ -324,12 +297,11 @@ class MMSingleStreamBlock(nn.Module): self.k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) # Fused operations with better naming - self.input_norm_scale_shift = LayerNormScaleShift( - hidden_size, - norm_type="layer", - eps=1e-6, - elementwise_affine=False, - dtype=dtype) + self.input_norm_scale_shift = LayerNormScaleShift(hidden_size, + norm_type="layer", + eps=1e-6, + elementwise_affine=False, + dtype=dtype) self.output_residual = ScaleResidual() # Activation function @@ -343,12 +315,11 @@ class MMSingleStreamBlock(nn.Module): prefix=f"{prefix}.modulation") # Distributed attention - self.attn = DistributedAttention( - num_heads=num_attention_heads, - head_size=head_dim, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.attn") + self.attn = DistributedAttention(num_heads=num_attention_heads, + head_size=head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn") def forward( self, @@ -368,9 +339,7 @@ class MMSingleStreamBlock(nn.Module): linear1_out, _ = self.linear1(x_mod) # Split into QKV and MLP parts - qkv, mlp = torch.split(linear1_out, - [3 * self.hidden_size, self.mlp_hidden_dim], - dim=-1) + qkv, mlp = torch.split(linear1_out, [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1) # Process QKV batch_size, seq_len = qkv.shape[0], qkv.shape[1] @@ -386,12 +355,16 @@ class MMSingleStreamBlock(nn.Module): img_k, txt_k = k[:, :-txt_len], k[:, -txt_len:] img_v, txt_v = v[:, :-txt_len], v[:, -txt_len:] - # Run distributed attention - img_attn_output, txt_attn_output = self.attn(img_q, img_k, img_v, original_seq_len, txt_q, - txt_k, txt_v, freqs_cis = freqs_cis) - attn_output = torch.cat((img_attn_output, txt_attn_output), - dim=1).view(batch_size, seq_len, -1) + img_attn_output, txt_attn_output = self.attn(img_q, + img_k, + img_v, + original_seq_len, + txt_q, + txt_k, + txt_v, + freqs_cis=freqs_cis) + attn_output = torch.cat((img_attn_output, txt_attn_output), dim=1).view(batch_size, seq_len, -1) # Process MLP activation mlp_output = self.mlp_act(mlp) @@ -421,19 +394,15 @@ class HunyuanVideoTransformer3DModel(BaseDiT): # shard single stream, double stream blocks, and refiner_blocks _fsdp_shard_conditions = HunyuanVideoConfig()._fsdp_shard_conditions _compile_conditions = HunyuanVideoConfig()._compile_conditions - _supported_attention_backends = HunyuanVideoConfig( - )._supported_attention_backends + _supported_attention_backends = HunyuanVideoConfig()._supported_attention_backends param_names_mapping = HunyuanVideoConfig().param_names_mapping - reverse_param_names_mapping = HunyuanVideoConfig( - ).reverse_param_names_mapping + reverse_param_names_mapping = HunyuanVideoConfig().reverse_param_names_mapping lora_param_names_mapping = HunyuanVideoConfig().lora_param_names_mapping def __init__(self, config: HunyuanVideoConfig, hf_config: dict[str, Any]): super().__init__(config=config, hf_config=hf_config) - self.patch_size = [ - config.patch_size_t, config.patch_size, config.patch_size - ] + self.patch_size = [config.patch_size_t, config.patch_size, config.patch_size] self.in_channels = config.in_channels self.num_channels_latents = config.num_channels_latents self.out_channels = config.in_channels if config.out_channels is None else config.out_channels @@ -448,9 +417,7 @@ class HunyuanVideoTransformer3DModel(BaseDiT): pe_dim = config.hidden_size // config.num_attention_heads if sum(config.rope_axes_dim) != pe_dim: - raise ValueError( - f"Got {config.rope_axes_dim} but expected positional dim {pe_dim}" - ) + raise ValueError(f"Got {config.rope_axes_dim} but expected positional dim {pe_dim}") self.hidden_size = config.hidden_size self.num_attention_heads = config.num_attention_heads @@ -486,33 +453,27 @@ class HunyuanVideoTransformer3DModel(BaseDiT): # Guidance modulation self.guidance_in = (TimestepEmbedder( - self.hidden_size, - act_layer="silu", - dtype=config.dtype, - prefix=f"{config.prefix}.guidance_in") + self.hidden_size, act_layer="silu", dtype=config.dtype, prefix=f"{config.prefix}.guidance_in") if self.guidance_embeds else None) # Double blocks self.double_blocks = nn.ModuleList([ - MMDoubleStreamBlock( - config.hidden_size, - config.num_attention_heads, - mlp_ratio=config.mlp_ratio, - dtype=config.dtype, - supported_attention_backends=self._supported_attention_backends, - prefix=f"{config.prefix}.double_blocks.{i}") - for i in range(config.num_layers) + MMDoubleStreamBlock(config.hidden_size, + config.num_attention_heads, + mlp_ratio=config.mlp_ratio, + dtype=config.dtype, + supported_attention_backends=self._supported_attention_backends, + prefix=f"{config.prefix}.double_blocks.{i}") for i in range(config.num_layers) ]) # Single blocks self.single_blocks = nn.ModuleList([ - MMSingleStreamBlock( - config.hidden_size, - config.num_attention_heads, - mlp_ratio=config.mlp_ratio, - dtype=config.dtype, - supported_attention_backends=self._supported_attention_backends, - prefix=f"{config.prefix}.single_blocks.{i+config.num_layers}") + MMSingleStreamBlock(config.hidden_size, + config.num_attention_heads, + mlp_ratio=config.mlp_ratio, + dtype=config.dtype, + supported_attention_backends=self._supported_attention_backends, + prefix=f"{config.prefix}.single_blocks.{i+config.num_layers}") for i in range(config.num_single_layers) ]) @@ -547,9 +508,7 @@ class HunyuanVideoTransformer3DModel(BaseDiT): Tuple of (output) """ if guidance is None: - guidance = torch.tensor([6016.0], - device=hidden_states.device, - dtype=hidden_states.dtype) + guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=hidden_states.dtype) img = x = hidden_states t = timestep @@ -571,9 +530,8 @@ class HunyuanVideoTransformer3DModel(BaseDiT): ) # Get rotary embeddings - freqs_cos, freqs_sin = get_rotary_pos_embed( - (tt, th, tw), self.hidden_size, - self.num_attention_heads, self.rope_dim_list, self.rope_theta) + freqs_cos, freqs_sin = get_rotary_pos_embed((tt, th, tw), self.hidden_size, self.num_attention_heads, + self.rope_dim_list, self.rope_theta) freqs_cos = freqs_cos.to(x.device) freqs_sin = freqs_sin.to(x.device) # Prepare modulation vectors @@ -644,18 +602,14 @@ class SingleTokenRefiner(nn.Module): super().__init__() # Input projection - self.input_embedder = ReplicatedLinear( - in_channels, - hidden_size, - bias=True, - params_dtype=dtype, - prefix=f"{prefix}.input_embedder") + self.input_embedder = ReplicatedLinear(in_channels, + hidden_size, + bias=True, + params_dtype=dtype, + prefix=f"{prefix}.input_embedder") # Timestep embedding - self.t_embedder = TimestepEmbedder(hidden_size, - act_layer="silu", - dtype=dtype, - prefix=f"{prefix}.t_embedder") + self.t_embedder = TimestepEmbedder(hidden_size, act_layer="silu", dtype=dtype, prefix=f"{prefix}.t_embedder") # Context embedding self.c_embedder = MLP(in_channels, @@ -684,8 +638,7 @@ class SingleTokenRefiner(nn.Module): context_aware_representations = torch.mean(x, dim=1) - context_aware_representations = self.c_embedder( - context_aware_representations) + context_aware_representations = self.c_embedder(context_aware_representations) c = timestep_aware_representations + context_aware_representations # Project input x, _ = self.input_embedder(x) @@ -714,10 +667,7 @@ class IndividualTokenRefinerBlock(nn.Module): mlp_hidden_dim = int(hidden_size * mlp_ratio) # Normalization and attention - self.norm1 = nn.LayerNorm(hidden_size, - eps=1e-6, - elementwise_affine=True, - dtype=dtype) + self.norm1 = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=True, dtype=dtype) self.self_attn_qkv = ReplicatedLinear(hidden_size, hidden_size * 3, @@ -725,40 +675,29 @@ class IndividualTokenRefinerBlock(nn.Module): params_dtype=dtype, prefix=f"{prefix}.self_attn_qkv") - self.self_attn_proj = ReplicatedLinear( - hidden_size, - hidden_size, - bias=qkv_bias, - params_dtype=dtype, - prefix=f"{prefix}.self_attn_proj") + self.self_attn_proj = ReplicatedLinear(hidden_size, + hidden_size, + bias=qkv_bias, + params_dtype=dtype, + prefix=f"{prefix}.self_attn_proj") # MLP - self.norm2 = nn.LayerNorm(hidden_size, - eps=1e-6, - elementwise_affine=True, - dtype=dtype) - self.mlp = MLP(hidden_size, - mlp_hidden_dim, - bias=True, - act_type="silu", - dtype=dtype, - prefix=f"{prefix}.mlp") + self.norm2 = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=True, dtype=dtype) + self.mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, act_type="silu", dtype=dtype, prefix=f"{prefix}.mlp") # Modulation - self.adaLN_modulation = ModulateProjection( - hidden_size, - factor=2, - act_layer="silu", - dtype=dtype, - prefix=f"{prefix}.adaLN_modulation") + self.adaLN_modulation = ModulateProjection(hidden_size, + factor=2, + act_layer="silu", + dtype=dtype, + prefix=f"{prefix}.adaLN_modulation") # Scaled dot product attention self.attn = LocalAttention( num_heads=num_attention_heads, head_size=hidden_size // num_attention_heads, # TODO: remove hardcode - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA), + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), ) def forward(self, x, c): @@ -774,8 +713,7 @@ class IndividualTokenRefinerBlock(nn.Module): # Run scaled dot product attention attn_output = self.attn(q, k, v) # [B, L, H, D] - attn_output = attn_output.reshape(batch_size, seq_len, - -1) # [B, L, H*D] + attn_output = attn_output.reshape(batch_size, seq_len, -1) # [B, L, H*D] # Project and apply residual connection with gating attn_out, _ = self.self_attn_proj(attn_output) @@ -793,19 +731,11 @@ class FinalLayer(nn.Module): The final layer of DiT that projects features to pixel space. """ - def __init__(self, - hidden_size, - patch_size, - out_channels, - dtype=None, - prefix: str = "") -> None: + def __init__(self, hidden_size, patch_size, out_channels, dtype=None, prefix: str = "") -> None: super().__init__() # Normalization - self.norm_final = nn.LayerNorm(hidden_size, - eps=1e-6, - elementwise_affine=False, - dtype=dtype) + self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False, dtype=dtype) output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels @@ -816,12 +746,11 @@ class FinalLayer(nn.Module): prefix=f"{prefix}.linear") # Modulation - self.adaLN_modulation = ModulateProjection( - hidden_size, - factor=2, - act_layer="silu", - dtype=dtype, - prefix=f"{prefix}.adaLN_modulation") + self.adaLN_modulation = ModulateProjection(hidden_size, + factor=2, + act_layer="silu", + dtype=dtype, + prefix=f"{prefix}.adaLN_modulation") def forward(self, x, c): # What the heck HF? Why you change the scale and shift order here??? @@ -830,5 +759,6 @@ class FinalLayer(nn.Module): x, _ = self.linear(x) return x + # Entry point for model registry EntryClass = HunyuanVideoTransformer3DModel diff --git a/fastvideo/models/dits/hunyuanvideo15.py b/fastvideo/models/dits/hunyuanvideo15.py index 7b32e9e21..a95fa8999 100644 --- a/fastvideo/models/dits/hunyuanvideo15.py +++ b/fastvideo/models/dits/hunyuanvideo15.py @@ -19,18 +19,15 @@ import torch.nn as nn from fastvideo.attention import DistributedAttention, LocalAttention from fastvideo.forward_context import set_forward_context -from fastvideo.distributed.communication_op import ( - sequence_model_parallel_all_gather_with_unpad, - sequence_model_parallel_shard) +from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad, + sequence_model_parallel_shard) from fastvideo.configs.models.dits import HunyuanVideo15Config -from fastvideo.layers.layernorm import (LayerNormScaleShift, ScaleResidual, - ScaleResidualLayerNormScaleShift) +from fastvideo.layers.layernorm import (LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift) from fastvideo.layers.linear import ReplicatedLinear # TODO(will-PY-refactor): RMSNorm .... from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import get_rotary_pos_embed -from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, - TimestepEmbedder, unpatchify) +from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, TimestepEmbedder, unpatchify) from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum from fastvideo.logger import init_logger @@ -39,6 +36,7 @@ from fastvideo.distributed.parallel_state import get_sp_world_size logger = init_logger(__name__) + class HunyuanRMSNorm(nn.Module): def __init__( @@ -135,6 +133,7 @@ class HunyuanVideo15TimeEmbedding(nn.Module): class HunyuanVideo15ByT5TextProjection(nn.Module): + def __init__(self, in_features: int, hidden_size: int, out_features: int): super().__init__() self.norm = nn.LayerNorm(in_features) @@ -154,6 +153,7 @@ class HunyuanVideo15ByT5TextProjection(nn.Module): class HunyuanVideo15ImageProjection(nn.Module): + def __init__(self, in_channels: int, hidden_size: int): super().__init__() self.norm_in = nn.LayerNorm(in_channels) @@ -204,15 +204,11 @@ class MMDoubleStreamBlock(nn.Module): ) # Fused operations for image stream - self.img_attn_norm = LayerNormScaleShift(hidden_size, - norm_type="layer", - elementwise_affine=False, - dtype=dtype) - self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift( - hidden_size, - norm_type="layer", - elementwise_affine=False, - dtype=dtype) + self.img_attn_norm = LayerNormScaleShift(hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype) + self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(hidden_size, + norm_type="layer", + elementwise_affine=False, + dtype=dtype) self.img_mlp_residual = ScaleResidual() # Image attention components @@ -231,11 +227,7 @@ class MMDoubleStreamBlock(nn.Module): params_dtype=dtype, prefix=f"{prefix}.img_attn_proj") - self.img_mlp = MLP(hidden_size, - mlp_hidden_dim, - bias=True, - dtype=dtype, - prefix=f"{prefix}.img_mlp") + self.img_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, prefix=f"{prefix}.img_mlp") # Text modulation components self.txt_mod = ModulateProjection( @@ -247,41 +239,30 @@ class MMDoubleStreamBlock(nn.Module): ) # Fused operations for text stream - self.txt_attn_norm = LayerNormScaleShift(hidden_size, - norm_type="layer", - elementwise_affine=False, - dtype=dtype) - self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift( - hidden_size, - norm_type="layer", - elementwise_affine=False, - dtype=dtype) + self.txt_attn_norm = LayerNormScaleShift(hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype) + self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(hidden_size, + norm_type="layer", + elementwise_affine=False, + dtype=dtype) self.txt_mlp_residual = ScaleResidual() # Text attention components - self.txt_attn_qkv = ReplicatedLinear(hidden_size, - hidden_size * 3, - bias=True, - params_dtype=dtype) + self.txt_attn_qkv = ReplicatedLinear(hidden_size, hidden_size * 3, bias=True, params_dtype=dtype) # QK norm layers for text self.txt_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) self.txt_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) - self.txt_attn_proj = ReplicatedLinear(hidden_size, - hidden_size, - bias=True, - params_dtype=dtype) + self.txt_attn_proj = ReplicatedLinear(hidden_size, hidden_size, bias=True, params_dtype=dtype) self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype) # Distributed attention - self.attn = DistributedAttention( - num_heads=num_attention_heads, - head_size=head_dim, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.attn") + self.attn = DistributedAttention(num_heads=num_attention_heads, + head_size=head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn") def forward( self, @@ -319,10 +300,8 @@ class MMDoubleStreamBlock(nn.Module): batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1] # Split QKV - img_qkv = img_qkv.view(batch_size, image_seq_len, 3, - self.num_attention_heads, -1) - img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, - 2] + img_qkv = img_qkv.view(batch_size, image_seq_len, 3, self.num_attention_heads, -1) + img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2] # Apply QK-Norm if needed img_q = self.img_attn_q_norm(img_q).to(img_v) @@ -336,34 +315,29 @@ class MMDoubleStreamBlock(nn.Module): batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1] # Split QKV - txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, - self.num_attention_heads, -1) - txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, - 2] + txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, self.num_attention_heads, -1) + txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2] # Apply QK-Norm if needed txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype) txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype) - img_attn, txt_attn = self.attn(img_q, img_k, img_v, original_seq_len, txt_q, txt_k, txt_v, freqs_cis=freqs_cis) - - img_attn_out, _ = self.img_attn_proj( - img_attn.view(batch_size, image_seq_len, -1)) + + img_attn_out, _ = self.img_attn_proj(img_attn.view(batch_size, image_seq_len, -1)) # Use fused operation for residual connection, normalization, and modulation - img_mlp_input, img_residual = self.img_attn_residual_mlp_norm( - img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale) + img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(img, img_attn_out, img_attn_gate, img_mlp_shift, + img_mlp_scale) # Process image MLP img_mlp_out = self.img_mlp(img_mlp_input) img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate) # Process text attention output - txt_attn_out, _ = self.txt_attn_proj( - txt_attn.reshape(batch_size, text_seq_len, -1)) + txt_attn_out, _ = self.txt_attn_proj(txt_attn.reshape(batch_size, text_seq_len, -1)) # Use fused operation for residual connection, normalization, and modulation - txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm( - txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale) + txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, + txt_mlp_scale) # Process text MLP txt_mlp_out = self.txt_mlp(txt_mlp_input) @@ -380,11 +354,9 @@ class HunyuanVideo15Transformer3DModel(BaseDiT): # shard single stream, double stream blocks, and refiner_blocks _fsdp_shard_conditions = HunyuanVideo15Config()._fsdp_shard_conditions _compile_conditions = HunyuanVideo15Config()._compile_conditions - _supported_attention_backends = HunyuanVideo15Config( - )._supported_attention_backends + _supported_attention_backends = HunyuanVideo15Config()._supported_attention_backends param_names_mapping = HunyuanVideo15Config().param_names_mapping - reverse_param_names_mapping = HunyuanVideo15Config( - ).reverse_param_names_mapping + reverse_param_names_mapping = HunyuanVideo15Config().reverse_param_names_mapping lora_param_names_mapping = HunyuanVideo15Config().lora_param_names_mapping def __init__( @@ -401,18 +373,18 @@ class HunyuanVideo15Transformer3DModel(BaseDiT): self.patch_size = (config.patch_size_t, config.patch_size, config.patch_size) # 1. Latent and condition embedders - self.img_in = PatchEmbed(self.patch_size, - config.in_channels, - self.hidden_size, - prefix=f"{config.prefix}.img_in") + self.img_in = PatchEmbed(self.patch_size, + config.in_channels, + self.hidden_size, + prefix=f"{config.prefix}.img_in") self.image_embedder = HunyuanVideo15ImageProjection(config.image_embed_dim, self.hidden_size) self.txt_in = SingleTokenRefiner(config.text_embed_dim, - self.hidden_size, - config.num_attention_heads, - depth=config.num_refiner_layers, - dtype=None, - prefix=f"{config.prefix}.txt_in") + self.hidden_size, + config.num_attention_heads, + depth=config.num_refiner_layers, + dtype=None, + prefix=f"{config.prefix}.txt_in") self.txt_in_2 = HunyuanVideo15ByT5TextProjection(config.text_embed_2_dim, 2048, self.hidden_size) @@ -422,25 +394,20 @@ class HunyuanVideo15Transformer3DModel(BaseDiT): # 3. Dual stream transformer blocks - self.double_blocks = nn.ModuleList( - [ - MMDoubleStreamBlock( - hidden_size=self.hidden_size, - num_attention_heads=config.num_attention_heads, - mlp_ratio=config.mlp_ratio, - dtype=None, - supported_attention_backends=self._supported_attention_backends, - prefix=f"{config.prefix}.double_blocks.{i}" - ) - for i in range(config.num_layers) - ] - ) + self.double_blocks = nn.ModuleList([ + MMDoubleStreamBlock(hidden_size=self.hidden_size, + num_attention_heads=config.num_attention_heads, + mlp_ratio=config.mlp_ratio, + dtype=None, + supported_attention_backends=self._supported_attention_backends, + prefix=f"{config.prefix}.double_blocks.{i}") for i in range(config.num_layers) + ]) # 5. Output projection self.final_layer = FinalLayer(self.hidden_size, - self.patch_size, - self.out_channels, - prefix=f"{config.prefix}.final_layer") + self.patch_size, + self.out_channels, + prefix=f"{config.prefix}.final_layer") self.gradient_checkpointing = False @@ -466,9 +433,9 @@ class HunyuanVideo15Transformer3DModel(BaseDiT): # 1. RoPE # Get rotary embeddings - freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames, post_patch_height, post_patch_width), self.hidden_size, - self.num_attention_heads, self.config.rope_axes_dim, self.config.rope_theta) + freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width), + self.hidden_size, self.num_attention_heads, + self.config.rope_axes_dim, self.config.rope_theta) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None @@ -482,59 +449,43 @@ class HunyuanVideo15Transformer3DModel(BaseDiT): current_seq_len = hidden_states.shape[1] sp_world_size = get_sp_world_size() padded_seq_len = current_seq_len * sp_world_size - - encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep) encoder_hidden_states_cond_emb = self.cond_type_embed( - torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long) - ) + torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long)) encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_cond_emb # byt5 text embedding encoder_hidden_states_2 = self.txt_in_2(encoder_hidden_states_2) encoder_hidden_states_2_cond_emb = self.cond_type_embed( - torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long) - ) + torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long)) encoder_hidden_states_2 = encoder_hidden_states_2 + encoder_hidden_states_2_cond_emb # image embed encoder_hidden_states_3 = self.image_embedder(encoder_hidden_states_image) is_t2v = torch.all(encoder_hidden_states_image == 0) - encoder_hidden_states_3_cond_emb = self.cond_type_embed( - 2 - * torch.ones_like( - encoder_hidden_states_3[:, :, 0], - dtype=torch.long, - ) - ) + encoder_hidden_states_3_cond_emb = self.cond_type_embed(2 * torch.ones_like( + encoder_hidden_states_3[:, :, 0], + dtype=torch.long, + )) encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb - encoder_hidden_states = torch.cat([encoder_hidden_states_2, encoder_hidden_states], dim=1) if is_t2v else torch.cat([encoder_hidden_states_3, encoder_hidden_states_2, encoder_hidden_states], dim=1) + encoder_hidden_states = torch.cat( + [encoder_hidden_states_2, encoder_hidden_states], dim=1) if is_t2v else torch.cat( + [encoder_hidden_states_3, encoder_hidden_states_2, encoder_hidden_states], dim=1) # 4. Transformer blocks if torch.is_grad_enabled() and self.gradient_checkpointing: for block in self.double_blocks: hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( - block, - hidden_states, - encoder_hidden_states, - temb, - freqs_cis, - original_seq_len - ) + block, hidden_states, encoder_hidden_states, temb, freqs_cis, original_seq_len) else: for block in self.double_blocks: - hidden_states, encoder_hidden_states = block( - hidden_states, - encoder_hidden_states, - temb, - freqs_cis, - original_seq_len - ) + hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, freqs_cis, + original_seq_len) # Final layer processing if get_sp_world_size() > 1: @@ -542,10 +493,12 @@ class HunyuanVideo15Transformer3DModel(BaseDiT): hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) hidden_states = self.final_layer(hidden_states, temb) # Unpatchify to get original shape - hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width, self.patch_size, self.out_channels) + hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width, + self.patch_size, self.out_channels) return hidden_states + class SingleTokenRefiner(nn.Module): """ A token refiner that processes text embeddings with attention to improve @@ -574,10 +527,7 @@ class SingleTokenRefiner(nn.Module): self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True) # Timestep embedding - self.t_embedder = TimestepEmbedder(hidden_size, - act_layer="silu", - dtype=dtype, - prefix=f"{prefix}.t_embedder") + self.t_embedder = TimestepEmbedder(hidden_size, act_layer="silu", dtype=dtype, prefix=f"{prefix}.t_embedder") # Context embedding self.c_embedder = MLP(in_channels, @@ -609,8 +559,7 @@ class SingleTokenRefiner(nn.Module): mask_float = mask.float().unsqueeze(-1) # [B, L, 1] context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1) - context_aware_representations = self.c_embedder( - context_aware_representations) + context_aware_representations = self.c_embedder(context_aware_representations) c = timestep_aware_representations + context_aware_representations # Project input x = self.input_embedder(x) @@ -619,6 +568,7 @@ class SingleTokenRefiner(nn.Module): x = block(x, c, mask=mask) return x + class IndividualTokenRefinerBlock(nn.Module): """ A transformer block for refining individual tokens with self-attention. @@ -638,10 +588,7 @@ class IndividualTokenRefinerBlock(nn.Module): mlp_hidden_dim = int(hidden_size * mlp_ratio) # Normalization and attention - self.norm1 = nn.LayerNorm(hidden_size, - eps=1e-6, - elementwise_affine=True, - dtype=dtype) + self.norm1 = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=True, dtype=dtype) self.self_attn_qkv = ReplicatedLinear(hidden_size, hidden_size * 3, @@ -649,40 +596,29 @@ class IndividualTokenRefinerBlock(nn.Module): params_dtype=dtype, prefix=f"{prefix}.self_attn_qkv") - self.self_attn_proj = ReplicatedLinear( - hidden_size, - hidden_size, - bias=qkv_bias, - params_dtype=dtype, - prefix=f"{prefix}.self_attn_proj") + self.self_attn_proj = ReplicatedLinear(hidden_size, + hidden_size, + bias=qkv_bias, + params_dtype=dtype, + prefix=f"{prefix}.self_attn_proj") # MLP - self.norm2 = nn.LayerNorm(hidden_size, - eps=1e-6, - elementwise_affine=True, - dtype=dtype) - self.mlp = MLP(hidden_size, - mlp_hidden_dim, - bias=True, - act_type="silu", - dtype=dtype, - prefix=f"{prefix}.mlp") + self.norm2 = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=True, dtype=dtype) + self.mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, act_type="silu", dtype=dtype, prefix=f"{prefix}.mlp") # Modulation - self.adaLN_modulation = ModulateProjection( - hidden_size, - factor=2, - act_layer="silu", - dtype=dtype, - prefix=f"{prefix}.adaLN_modulation") + self.adaLN_modulation = ModulateProjection(hidden_size, + factor=2, + act_layer="silu", + dtype=dtype, + prefix=f"{prefix}.adaLN_modulation") # Scaled dot product attention self.attn = LocalAttention( num_heads=num_attention_heads, head_size=hidden_size // num_attention_heads, # TODO: remove hardcode - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA), + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), ) def forward(self, x, c, mask=None): @@ -697,7 +633,7 @@ class IndividualTokenRefinerBlock(nn.Module): qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] if mask is None: - attn_output = self.attn(q, k, v) + attn_output = self.attn(q, k, v) else: from fastvideo.attention.backends.sdpa import SDPAMetadataBuilder attn_metadata = SDPAMetadataBuilder().build( @@ -707,8 +643,7 @@ class IndividualTokenRefinerBlock(nn.Module): # Run distributed attention with set_forward_context(current_timestep=0, attn_metadata=attn_metadata): attn_output = self.attn(q, k, v) # [B, L, H, D] - attn_output = attn_output.reshape(batch_size, seq_len, - -1) + attn_output = attn_output.reshape(batch_size, seq_len, -1) # Project and apply residual connection with gating attn_out, _ = self.self_attn_proj(attn_output) @@ -720,24 +655,17 @@ class IndividualTokenRefinerBlock(nn.Module): return x + class FinalLayer(nn.Module): """ The final layer of DiT that projects features to pixel space. """ - def __init__(self, - hidden_size, - patch_size, - out_channels, - dtype=None, - prefix: str = "") -> None: + def __init__(self, hidden_size, patch_size, out_channels, dtype=None, prefix: str = "") -> None: super().__init__() # Normalization - self.norm_final = nn.LayerNorm(hidden_size, - eps=1e-6, - elementwise_affine=False, - dtype=dtype) + self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False, dtype=dtype) output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels @@ -748,12 +676,11 @@ class FinalLayer(nn.Module): prefix=f"{prefix}.linear") # Modulation - self.adaLN_modulation = ModulateProjection( - hidden_size, - factor=2, - act_layer="silu", - dtype=dtype, - prefix=f"{prefix}.adaLN_modulation") + self.adaLN_modulation = ModulateProjection(hidden_size, + factor=2, + act_layer="silu", + dtype=dtype, + prefix=f"{prefix}.adaLN_modulation") def forward(self, x, c): # What the heck HF? Why you change the scale and shift order here??? @@ -762,5 +689,6 @@ class FinalLayer(nn.Module): x, _ = self.linear(x) return x + # Entry point for model registry EntryClass = HunyuanVideo15Transformer3DModel diff --git a/fastvideo/models/dits/hyworld/__init__.py b/fastvideo/models/dits/hyworld/__init__.py index 03516980a..5213bf913 100644 --- a/fastvideo/models/dits/hyworld/__init__.py +++ b/fastvideo/models/dits/hyworld/__init__.py @@ -12,8 +12,7 @@ from .hyworld import HYWorldTransformer3DModel, HYWorldDoubleStreamBlock # Inference utilities (used by examples) from .resolution_utils import ( - get_resolution_from_image, -) + get_resolution_from_image, ) __all__ = [ # Model (used by model registry) diff --git a/fastvideo/models/dits/hyworld/camera_rope.py b/fastvideo/models/dits/hyworld/camera_rope.py index 672988ef0..8f1463e9c 100644 --- a/fastvideo/models/dits/hyworld/camera_rope.py +++ b/fastvideo/models/dits/hyworld/camera_rope.py @@ -126,9 +126,9 @@ def _prepare_apply_fns_all_dim( coeffs_x: Optional[torch.Tensor] = None, coeffs_y: Optional[torch.Tensor] = None, ) -> Tuple[ - Callable[[torch.Tensor], torch.Tensor], - Callable[[torch.Tensor], torch.Tensor], - Callable[[torch.Tensor], torch.Tensor], + Callable[[torch.Tensor], torch.Tensor], + Callable[[torch.Tensor], torch.Tensor], + Callable[[torch.Tensor], torch.Tensor], ]: """Prepare transforms for PRoPE-style positional encoding.""" device = viewmats.device @@ -184,8 +184,8 @@ def _prepare_apply_fns_all_dim( def _apply_tiled_projmat( - feats: torch.Tensor, # (batch, num_heads, seqlen, feat_dim) - matrix: torch.Tensor, # (batch, cameras, D, D) + feats: torch.Tensor, # (batch, num_heads, seqlen, feat_dim) + matrix: torch.Tensor, # (batch, cameras, D, D) ) -> torch.Tensor: """Apply projection matrix to features.""" # - seqlen => (cameras, patches_x * patches_y) diff --git a/fastvideo/models/dits/hyworld/data_utils.py b/fastvideo/models/dits/hyworld/data_utils.py index f5d1b5454..a12575a90 100644 --- a/fastvideo/models/dits/hyworld/data_utils.py +++ b/fastvideo/models/dits/hyworld/data_utils.py @@ -62,7 +62,7 @@ def get_closest_ratio(height: float, width: float, ratios: list, buckets: list): def generate_crop_size_list(base_size=256, patch_size=16, max_ratio=4.0): - num_patches = round((base_size / patch_size) ** 2) + num_patches = round((base_size / patch_size)**2) assert max_ratio >= 1.0 crop_size_list = [] wp, hp = num_patches, 1 diff --git a/fastvideo/models/dits/hyworld/hyworld.py b/fastvideo/models/dits/hyworld/hyworld.py index b16cd72e5..91534d1af 100644 --- a/fastvideo/models/dits/hyworld/hyworld.py +++ b/fastvideo/models/dits/hyworld/hyworld.py @@ -20,9 +20,8 @@ import torch import torch.nn as nn from einops import rearrange, repeat -from fastvideo.distributed.communication_op import ( - sequence_model_parallel_all_gather_with_unpad, - sequence_model_parallel_shard) +from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad, + sequence_model_parallel_shard) from fastvideo.configs.models.dits import HYWorldConfig from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.rotary_embedding import get_rotary_pos_embed @@ -67,13 +66,11 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): self.hidden_size = hidden_size # Add ProPE projection layer for camera-aware attention - self.img_attn_prope_proj = ReplicatedLinear( - hidden_size, - hidden_size, - bias=True, - params_dtype=dtype, - prefix=f"{prefix}.img_attn_prope_proj" - ) + self.img_attn_prope_proj = ReplicatedLinear(hidden_size, + hidden_size, + bias=True, + params_dtype=dtype, + prefix=f"{prefix}.img_attn_prope_proj") # Zero-initialize ProPE projection (starts as identity) nn.init.zeros_(self.img_attn_prope_proj.weight) if self.img_attn_prope_proj.bias is not None: @@ -135,10 +132,8 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1] # Split QKV - img_qkv = img_qkv.view(batch_size, image_seq_len, 3, - self.num_attention_heads, -1) - img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, - 2] + img_qkv = img_qkv.view(batch_size, image_seq_len, 3, self.num_attention_heads, -1) + img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2] # Apply QK-Norm if needed img_q = self.img_attn_q_norm(img_q).to(img_v) @@ -152,10 +147,8 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1] # Split QKV - txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, - self.num_attention_heads, -1) - txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, - 2] + txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, self.num_attention_heads, -1) + txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2] # Apply QK-Norm if needed txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype) txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype) @@ -193,7 +186,7 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): txt_v, freqs_cis=freqs_cis, ) - + # begin hyworld # attention with prope (same text mask, so reuse the metadata) # NOTE: Do NOT pass freqs_cis to prope attention - HY-WorldPlay does not apply RoPE to prope @@ -209,10 +202,8 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): freqs_cis=None, # No RoPE for prope attention ) img_attn_prope = img_attn_prope.reshape(batch_size, image_seq_len, -1) - img_attn_prope = rearrange( - img_attn_prope, "B L (H D) -> B H L D", H=self.num_attention_heads - ) - img_attn_prope = apply_fn_o(img_attn_prope) # [batch, num_heads, seqlen, head_dim] + img_attn_prope = rearrange(img_attn_prope, "B L (H D) -> B H L D", H=self.num_attention_heads) + img_attn_prope = apply_fn_o(img_attn_prope) # [batch, num_heads, seqlen, head_dim] img_attn_prope = rearrange(img_attn_prope, "B H L D -> B L (H D)") # add prope to img_attn @@ -222,20 +213,27 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): # end hyworld # Use fused operation for residual connection, normalization, and modulation - img_mlp_input, img_residual = self.img_attn_residual_mlp_norm( - img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale, convert_modulation_dtype=True) + img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(img, + img_attn_out, + img_attn_gate, + img_mlp_shift, + img_mlp_scale, + convert_modulation_dtype=True) # Process image MLP img_mlp_out = self.img_mlp(img_mlp_input) img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate) # Process text attention output - txt_attn_out, _ = self.txt_attn_proj( - txt_attn.reshape(batch_size, text_seq_len, -1)) + txt_attn_out, _ = self.txt_attn_proj(txt_attn.reshape(batch_size, text_seq_len, -1)) # Use fused operation for residual connection, normalization, and modulation - txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm( - txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale, convert_modulation_dtype=True) + txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(txt, + txt_attn_out, + txt_attn_gate, + txt_mlp_shift, + txt_mlp_scale, + convert_modulation_dtype=True) # Process text MLP txt_mlp_out = self.txt_mlp(txt_mlp_input) @@ -250,24 +248,18 @@ class HYWorldFinalLayer(nn.Module): This matches HY-WorldPlay's FinalLayer behavior. """ - def __init__(self, - hidden_size, - patch_size, - out_channels, - dtype=None, - prefix: str = "") -> None: + def __init__(self, hidden_size, patch_size, out_channels, dtype=None, prefix: str = "") -> None: super().__init__() from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.visual_embedding import ModulateProjection from fastvideo.layers.layernorm import LayerNormScaleShift - self.norm_final = LayerNormScaleShift( - hidden_size, - norm_type="layer", - eps=1e-6, - elementwise_affine=False, - dtype=dtype, - prefix=f"{prefix}.norm_final") + self.norm_final = LayerNormScaleShift(hidden_size, + norm_type="layer", + eps=1e-6, + elementwise_affine=False, + dtype=dtype, + prefix=f"{prefix}.norm_final") output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels @@ -278,12 +270,11 @@ class HYWorldFinalLayer(nn.Module): prefix=f"{prefix}.linear") # Modulation projection to get shift/scale - self.adaLN_modulation = ModulateProjection( - hidden_size, - factor=2, - act_layer="silu", - dtype=dtype, - prefix=f"{prefix}.adaLN_modulation") + self.adaLN_modulation = ModulateProjection(hidden_size, + factor=2, + act_layer="silu", + dtype=dtype, + prefix=f"{prefix}.adaLN_modulation") def forward(self, x, c): shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1) @@ -298,7 +289,7 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): - ProPE (Projective Positional Encoding) for camera-aware attention - Action conditioning for interactive video generation """ - + # Class attributes for weight loading - use HYWorld-specific mapping _fsdp_shard_conditions = HYWorldConfig().arch_config._fsdp_shard_conditions _compile_conditions = HYWorldConfig().arch_config._compile_conditions @@ -314,37 +305,31 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): # Replace double_blocks with HY-World version that supports ProPE self.double_blocks = nn.ModuleList([ - HYWorldDoubleStreamBlock( - hidden_size=self.hidden_size, - num_attention_heads=self.num_attention_heads, - mlp_ratio=config.arch_config.mlp_ratio, - dtype=None, - supported_attention_backends=self._supported_attention_backends, - prefix=f"{config.prefix}.double_blocks.{i}" - ) + HYWorldDoubleStreamBlock(hidden_size=self.hidden_size, + num_attention_heads=self.num_attention_heads, + mlp_ratio=config.arch_config.mlp_ratio, + dtype=None, + supported_attention_backends=self._supported_attention_backends, + prefix=f"{config.prefix}.double_blocks.{i}") for i in range(config.arch_config.num_layers) ]) # Add action conditioning module - self.action_in = TimestepEmbedder( - self.hidden_size, - act_layer="silu", - dtype=None, - prefix=f"{config.prefix}.action_in" - ) + self.action_in = TimestepEmbedder(self.hidden_size, + act_layer="silu", + dtype=None, + prefix=f"{config.prefix}.action_in") # Zero-initialize action embedding (starts with no effect) nn.init.zeros_(self.action_in.mlp.fc_out.weight) if self.action_in.mlp.fc_out.bias is not None: nn.init.zeros_(self.action_in.mlp.fc_out.bias) # Override final_layer with HYWorld version that uses per-token modulate() - self.final_layer = HYWorldFinalLayer( - hidden_size=self.hidden_size, - patch_size=self.patch_size, - out_channels=self.out_channels, - dtype=None, - prefix=f"{config.prefix}.final_layer" - ) + self.final_layer = HYWorldFinalLayer(hidden_size=self.hidden_size, + patch_size=self.patch_size, + out_channels=self.out_channels, + dtype=None, + prefix=f"{config.prefix}.final_layer") self.gradient_checkpointing = False @@ -384,13 +369,9 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): # 1. RoPE # Get rotary embeddings - freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames, post_patch_height, post_patch_width), - self.hidden_size, - self.num_attention_heads, - self.config.rope_axes_dim, - self.config.rope_theta - ) + freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width), + self.hidden_size, self.num_attention_heads, + self.config.rope_axes_dim, self.config.rope_theta) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) # NOTE: freqs_cis does NOT need sharding because FastVideo's DistributedAttention @@ -404,7 +385,7 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): # Add action conditioning if provided # temb shape: [B*T, C] where T = num_frames temb = temb + self.action_in(action.reshape(-1)) - + # Broadcast timestep embedding for transformer blocks (one per spatial token) # [B*T, C] -> [B, T*H*W, C] -> [B*T*H*W, C] temb = repeat(temb, "(B T) C -> B (T H W) C", B=batch_size, H=post_patch_height, W=post_patch_width) @@ -413,17 +394,9 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1) sp_world_size = get_sp_world_size() - viewmats_seq = repeat( - viewmats, "B T M N->B (T H W) M N", - H=post_patch_height, - W=post_patch_width - ) - Ks_seq = repeat( - Ks, "B T M N->B (T H W) M N", - H=post_patch_height, - W=post_patch_width - ) - + viewmats_seq = repeat(viewmats, "B T M N->B (T H W) M N", H=post_patch_height, W=post_patch_width) + Ks_seq = repeat(Ks, "B T M N->B (T H W) M N", H=post_patch_height, W=post_patch_width) + # Shard viewmats, Ks, and temb for sequence parallelism (shard along sequence dim=1) # Note that temb in HY1.5 does not need sharding because it is per-sample modulation # In HYWorld, temb is per-token modulation. @@ -431,7 +404,7 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): viewmats_seq, _ = sequence_model_parallel_shard(viewmats_seq, dim=1) Ks_seq, _ = sequence_model_parallel_shard(Ks_seq, dim=1) temb, _ = sequence_model_parallel_shard(temb, dim=1) - + # Rearrange temb after sharding to match expected shape temb = rearrange(temb, "B S C -> (B S) C") @@ -439,16 +412,14 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep_txt, encoder_attention_mask) encoder_hidden_states_cond_emb = self.cond_type_embed( - torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long) - ) + torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long)) encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_cond_emb # byt5 text embedding encoder_hidden_states_2 = self.txt_in_2(encoder_hidden_states_2) encoder_hidden_states_2_cond_emb = self.cond_type_embed( - torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long) - ) + torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long)) encoder_hidden_states_2 = encoder_hidden_states_2 + encoder_hidden_states_2_cond_emb # image embed @@ -467,13 +438,10 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): dtype=encoder_attention_mask.dtype, device=encoder_attention_mask.device, ) - encoder_hidden_states_3_cond_emb = self.cond_type_embed( - 2 - * torch.ones_like( - encoder_hidden_states_3[:, :, 0], - dtype=torch.long, - ) - ) + encoder_hidden_states_3_cond_emb = self.cond_type_embed(2 * torch.ones_like( + encoder_hidden_states_3[:, :, 0], + dtype=torch.long, + )) encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb # reorder and combine text tokens: combine valid tokens first, then padding @@ -484,12 +452,12 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): new_encoder_attention_mask = [] for text, text_mask, text_2, text_mask_2, image, image_mask in zip( - encoder_hidden_states, - encoder_attention_mask, - encoder_hidden_states_2, - encoder_attention_mask_2, - encoder_hidden_states_3, - encoder_attention_mask_3, + encoder_hidden_states, + encoder_attention_mask, + encoder_hidden_states_2, + encoder_attention_mask_2, + encoder_hidden_states_3, + encoder_attention_mask_3, ): # Concatenate: [valid_image, valid_byt5, valid_mllm, invalid_image, invalid_byt5, invalid_mllm] new_encoder_hidden_states.append( @@ -503,8 +471,7 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): torch.zeros_like(text[~text_mask]), # invalid mllm (zeroed) ], dim=0, - ) - ) + )) # Apply same reordering to attention masks new_encoder_attention_mask.append( torch.cat( @@ -517,8 +484,7 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): text_mask[~text_mask], ], dim=0, - ) - ) + )) encoder_hidden_states = torch.stack(new_encoder_hidden_states) encoder_attention_mask = torch.stack(new_encoder_attention_mask) @@ -552,17 +518,17 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): Ks_seq, ) - # Final layer processing (per-token conditioning via HYWorldFinalLayer) # Apply final_layer on sharded data first, then gather hidden_states = self.final_layer(hidden_states, temb) - + # Gather the output from all ranks if get_sp_world_size() > 1: hidden_states = hidden_states.contiguous() hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) - + # Unpatchify to get original shape - hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width, self.patch_size, self.out_channels) + hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width, + self.patch_size, self.out_channels) return hidden_states diff --git a/fastvideo/models/dits/hyworld/pose.py b/fastvideo/models/dits/hyworld/pose.py index b1b3f5df3..40410591f 100644 --- a/fastvideo/models/dits/hyworld/pose.py +++ b/fastvideo/models/dits/hyworld/pose.py @@ -1,5 +1,4 @@ # Some functions from HY-WorldPlay/hyvideo/generate.py - """ Pose processing utilities for HYWorld video generation. @@ -17,7 +16,6 @@ from typing import Union, Optional from .trajectory import generate_camera_trajectory_local - # Mapping from one-hot action encoding to single label mapping = { (0, 0, 0, 0): 0, @@ -87,9 +85,7 @@ def parse_pose_string( parts = cmd.split("-") if len(parts) != 2: - raise ValueError( - f"Invalid pose command: {cmd}. Expected format: 'action-duration'" - ) + raise ValueError(f"Invalid pose command: {cmd}. Expected format: 'action-duration'") action = parts[0].strip() try: @@ -133,13 +129,12 @@ def parse_pose_string( for _ in range(num_frames): motions.append({"yaw": yaw_speed}) else: - raise ValueError( - f"Unknown action: {action}. " - f"Supported actions: w, s, a, d, up, down, left, right" - ) + raise ValueError(f"Unknown action: {action}. " + f"Supported actions: w, s, a, d, up, down, left, right") return motions + def pose_string_to_json( pose_string: str, intrinsic: Optional[list[list[float]]] = None, @@ -156,7 +151,7 @@ def pose_string_to_json( """ if intrinsic is None: intrinsic = DEFAULT_INTRINSIC - + motions = parse_pose_string(pose_string) poses = generate_camera_trajectory_local(motions) @@ -166,6 +161,7 @@ def pose_string_to_json( return pose_json + def pose_to_input( pose_data: Union[str, dict], latent_num: int, @@ -200,16 +196,13 @@ def pose_to_input( elif isinstance(pose_data, dict): pose_json = pose_data else: - raise ValueError( - f"Invalid pose_data type: {type(pose_data)}. Expected str or dict." - ) + raise ValueError(f"Invalid pose_data type: {type(pose_data)}. Expected str or dict.") pose_keys = list(pose_json.keys()) latent_num_from_pose = len(pose_keys) assert latent_num_from_pose == latent_num, ( f"pose corresponds to {latent_num_from_pose * 4 - 3} frames, num_frames " - f"must be set to {latent_num_from_pose * 4 - 3} to ensure alignment." - ) + f"must be set to {latent_num_from_pose * 4 - 3} to ensure alignment.") intrinsic_list = [] w2c_list = [] @@ -218,7 +211,7 @@ def pose_to_input( c2w = np.array(pose_json[t_key]["extrinsic"]) w2c = np.linalg.inv(c2w) w2c_list.append(w2c) - + # Normalize intrinsics intrinsic = np.array(pose_json[t_key]["K"]) intrinsic[0, 0] /= intrinsic[0, 2] * 2 @@ -236,7 +229,7 @@ def pose_to_input( relative_c2w = np.zeros_like(c2ws) relative_c2w[0, ...] = c2ws[0, ...] relative_c2w[1:, ...] = C_inv @ c2ws[1:, ...] - + # Initialize one-hot action encodings trans_one_hot = np.zeros((relative_c2w.shape[0], 4), dtype=np.int32) rotate_one_hot = np.zeros((relative_c2w.shape[0], 4), dtype=np.int32) @@ -245,7 +238,7 @@ def pose_to_input( for i in range(1, relative_c2w.shape[0]): move_dirs = relative_c2w[i, :3, 3] # direction vector move_norms = np.linalg.norm(move_dirs) - + if move_norms > move_norm_valid: # threshold for movement move_norm_dirs = move_dirs / move_norms angles_rad = np.arccos(move_norm_dirs.clip(-1.0, 1.0)) @@ -259,9 +252,7 @@ def pose_to_input( # Determine movement and rotation actions if move_norms > move_norm_valid: # threshold for movement - if (not tps) or ( - tps and abs(rot_angles_deg[1]) < 5e-2 and abs(rot_angles_deg[0]) < 5e-2 - ): + if (not tps) or (tps and abs(rot_angles_deg[1]) < 5e-2 and abs(rot_angles_deg[0]) < 5e-2): if trans_angles_deg[2] < 60: trans_one_hot[i, 0] = 1 # forward elif trans_angles_deg[2] > 120: @@ -305,7 +296,6 @@ def camera_center_normalization(w2c: np.ndarray) -> np.ndarray: return np.linalg.inv(c2w_aligned) - def parse_pose_string_to_actions(pose_string: str, fps: int = 24) -> list[dict]: """ Parse pose string to frame-level action timeline. @@ -335,9 +325,7 @@ def parse_pose_string_to_actions(pose_string: str, fps: int = 24) -> list[dict]: parts = cmd.split("-") if len(parts) != 2: - raise ValueError( - f"Invalid pose command: {cmd}. Expected format: 'action-duration'" - ) + raise ValueError(f"Invalid pose command: {cmd}. Expected format: 'action-duration'") action = parts[0].strip() try: diff --git a/fastvideo/models/dits/hyworld/resolution_utils.py b/fastvideo/models/dits/hyworld/resolution_utils.py index 7797c12a8..737c45f81 100644 --- a/fastvideo/models/dits/hyworld/resolution_utils.py +++ b/fastvideo/models/dits/hyworld/resolution_utils.py @@ -7,10 +7,22 @@ from fastvideo.models.dits.hyworld.data_utils import generate_crop_size_list # Target resolution configs (matching HY-WorldPlay) TARGET_SIZE_CONFIG = { - "360p": {"bucket_hw_base_size": 480, "bucket_hw_bucket_stride": 16}, - "480p": {"bucket_hw_base_size": 640, "bucket_hw_bucket_stride": 16}, - "720p": {"bucket_hw_base_size": 960, "bucket_hw_bucket_stride": 16}, - "1080p": {"bucket_hw_base_size": 1440, "bucket_hw_bucket_stride": 16}, + "360p": { + "bucket_hw_base_size": 480, + "bucket_hw_bucket_stride": 16 + }, + "480p": { + "bucket_hw_base_size": 640, + "bucket_hw_bucket_stride": 16 + }, + "720p": { + "bucket_hw_base_size": 960, + "bucket_hw_bucket_stride": 16 + }, + "1080p": { + "bucket_hw_base_size": 1440, + "bucket_hw_bucket_stride": 16 + }, } @@ -29,15 +41,15 @@ def get_closest_resolution(image_height, image_width, target_resolution="480p"): config = TARGET_SIZE_CONFIG[target_resolution] bucket_hw_base_size = config["bucket_hw_base_size"] bucket_hw_bucket_stride = config["bucket_hw_bucket_stride"] - + crop_size_list = generate_crop_size_list(bucket_hw_base_size, bucket_hw_bucket_stride) aspect_ratios = np.array([round(float(h) / float(w), 5) for h, w in crop_size_list]) - + # Find closest aspect ratio image_ratio = float(image_height) / float(image_width) closest_idx = np.abs(aspect_ratios - image_ratio).argmin() closest_size = crop_size_list[closest_idx] - + return closest_size[0], closest_size[1] # (height, width) @@ -61,4 +73,3 @@ def get_resolution_from_image(image_path, target_resolution="480p"): img = Image.open(image_path) img_width, img_height = img.size return get_closest_resolution(img_height, img_width, target_resolution) - diff --git a/fastvideo/models/dits/hyworld/retrieval_context.py b/fastvideo/models/dits/hyworld/retrieval_context.py index f50813bd5..e5ccf3343 100644 --- a/fastvideo/models/dits/hyworld/retrieval_context.py +++ b/fastvideo/models/dits/hyworld/retrieval_context.py @@ -243,10 +243,8 @@ def select_aligned_memory_frames( num_total_frames = len(w2c_list) if current_frame_idx >= num_total_frames or current_frame_idx < 3: - raise ValueError( - f"The current frame index must be within the valid range of w2c_list and must be at least 3." - f"{current_frame_idx}, {len(w2c_list)}" - ) + raise ValueError(f"The current frame index must be within the valid range of w2c_list and must be at least 3." + f"{current_frame_idx}, {len(w2c_list)}") start_context_idx = max(0, current_frame_idx - temporal_context_size) context_frames_indices = list(range(start_context_idx, current_frame_idx)) @@ -255,17 +253,11 @@ def select_aligned_memory_frames( query_clip_indices = list( range( current_frame_idx, - ( - current_frame_idx + pred_latent_size - if current_frame_idx + pred_latent_size <= num_total_frames - else num_total_frames - ), - ) - ) + (current_frame_idx + pred_latent_size if current_frame_idx + + pred_latent_size <= num_total_frames else num_total_frames), + )) - historical_clip_indices = list( - range(4, current_frame_idx - temporal_context_size, 4) - ) + historical_clip_indices = list(range(4, current_frame_idx - temporal_context_size, 4)) memory_frames_indices = [0, 1, 2, 3] # add the first chunk as context memory_frames = memory_frames - temporal_context_size diff --git a/fastvideo/models/dits/hyworld/trajectory.py b/fastvideo/models/dits/hyworld/trajectory.py index 6fccfb03e..5d7699b0b 100644 --- a/fastvideo/models/dits/hyworld/trajectory.py +++ b/fastvideo/models/dits/hyworld/trajectory.py @@ -66,13 +66,11 @@ def generate_camera_trajectory_local(motions): C = np.array([[1, 0.0, 0, 0], [0, 1, 0, 0], [0, 0, 1, -1.0], [0, 0, 0, 1]]) c_origin = C.copy() # Rotation around the Y-axis - R_y = np.array( - [ - [np.cos(theta), 0, np.sin(theta)], - [0, 1, 0], - [-np.sin(theta), 0, np.cos(theta)], - ] - ) + R_y = np.array([ + [np.cos(theta), 0, np.sin(theta)], + [0, 1, 0], + [-np.sin(theta), 0, np.cos(theta)], + ]) # Translation C[:3, :3] = C[:3, :3] @ R_y C[:3, 3] = R_y @ C[:3, 3] diff --git a/fastvideo/models/dits/kandinsky5.py b/fastvideo/models/dits/kandinsky5.py index 025c32482..e0365c511 100644 --- a/fastvideo/models/dits/kandinsky5.py +++ b/fastvideo/models/dits/kandinsky5.py @@ -31,11 +31,11 @@ _ARCH_CONFIG_DEFAULTS = Kandinsky5VideoConfig().arch_config def _build_rotary_freqs(dim: int, max_period: float) -> torch.Tensor: - return torch.exp(-math.log(max_period) * torch.arange( - start=0, end=dim, dtype=torch.float32) / dim) + return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) -def local_patching(x: torch.Tensor, shape: tuple[int, int, int, int], +def local_patching(x: torch.Tensor, + shape: tuple[int, int, int, int], group_size: tuple[int, int, int], dim: int = 0) -> torch.Tensor: batch_size, duration, height, width = shape @@ -64,7 +64,8 @@ def local_patching(x: torch.Tensor, shape: tuple[int, int, int, int], return x -def local_merge(x: torch.Tensor, shape: tuple[int, int, int, int], +def local_merge(x: torch.Tensor, + shape: tuple[int, int, int, int], group_size: tuple[int, int, int], dim: int = 0) -> torch.Tensor: batch_size, duration, height, width = shape @@ -89,15 +90,11 @@ def local_merge(x: torch.Tensor, shape: tuple[int, int, int, int], dim + 5, *range(dim + 6, len(x.shape)), ) - x = x.flatten(dim, dim + 1).flatten(dim + 1, dim + 2).flatten( - dim + 2, dim + 3) + x = x.flatten(dim, dim + 1).flatten(dim + 1, dim + 2).flatten(dim + 2, dim + 3) return x -def fractal_flatten(x: torch.Tensor, - rope: torch.Tensor, - shape: tuple[int, int, int, int], - block_mask: bool = False): +def fractal_flatten(x: torch.Tensor, rope: torch.Tensor, shape: tuple[int, int, int, int], block_mask: bool = False): if block_mask: pixel_size = FRACTAL_PIXEL_SIZE x = local_patching(x, shape, (1, pixel_size, pixel_size), dim=1) @@ -110,9 +107,7 @@ def fractal_flatten(x: torch.Tensor, return x, rope -def fractal_unflatten(x: torch.Tensor, - shape: tuple[int, int, int, int], - block_mask: bool = False): +def fractal_unflatten(x: torch.Tensor, shape: tuple[int, int, int, int], block_mask: bool = False): if block_mask: pixel_size = FRACTAL_PIXEL_SIZE x = x.reshape(x.shape[0], -1, pixel_size**2, *x.shape[2:]) @@ -124,8 +119,7 @@ def fractal_unflatten(x: torch.Tensor, class Kandinsky5TimeEmbeddings(nn.Module): - def __init__(self, model_dim: int, time_dim: int, - max_period: float = 10000.0): + def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): super().__init__() assert model_dim % 2 == 0 self.model_dim = model_dim @@ -159,12 +153,10 @@ class Kandinsky5TextEmbeddings(nn.Module): class Kandinsky5VisualEmbeddings(nn.Module): - def __init__(self, visual_dim: int, model_dim: int, - patch_size: tuple[int, int, int]): + def __init__(self, visual_dim: int, model_dim: int, patch_size: tuple[int, int, int]): super().__init__() self.patch_size = patch_size - self.in_layer = ReplicatedLinear(math.prod(patch_size) * visual_dim, - model_dim) + self.in_layer = ReplicatedLinear(math.prod(patch_size) * visual_dim, model_dim) def forward(self, x: torch.Tensor) -> torch.Tensor: batch_size, duration, height, width, dim = x.shape @@ -184,10 +176,7 @@ class Kandinsky5VisualEmbeddings(nn.Module): class Kandinsky5RoPE1D(nn.Module): - def __init__(self, - dim: int, - max_pos: int = 1024, - max_period: float = 10000.0): + def __init__(self, dim: int, max_pos: int = 1024, max_period: float = 10000.0): super().__init__() self.max_period = max_period self.dim = dim @@ -216,20 +205,16 @@ class Kandinsky5RoPE3D(nn.Module): self.max_pos = max_pos self.max_period = max_period - for i, (axes_dim, ax_max_pos) in enumerate( - zip(axes_dims, max_pos, strict=True)): + for i, (axes_dim, ax_max_pos) in enumerate(zip(axes_dims, max_pos, strict=True)): freq = _build_rotary_freqs(axes_dim // 2, max_period) pos = torch.arange(ax_max_pos, dtype=freq.dtype) - self.register_buffer(f"args_{i}", - torch.outer(pos, freq), - persistent=False) + self.register_buffer(f"args_{i}", torch.outer(pos, freq), persistent=False) def forward(self, shape: tuple[int, int, int, int], pos: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | list[torch.Tensor], - scale_factor: tuple[float, float, - float] = (1.0, 1.0, 1.0)): + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0)): batch_size, duration, height, width = shape args_t = self.args_0[pos[0]] / scale_factor[0] args_h = self.args_1[pos[1]] / scale_factor[1] @@ -237,12 +222,9 @@ class Kandinsky5RoPE3D(nn.Module): args = torch.cat( [ - args_t.view(1, duration, 1, 1, -1).repeat( - batch_size, 1, height, width, 1), - args_h.view(1, 1, height, 1, -1).repeat( - batch_size, duration, 1, width, 1), - args_w.view(1, 1, 1, width, -1).repeat( - batch_size, duration, height, 1, 1), + args_t.view(1, duration, 1, 1, -1).repeat(batch_size, 1, height, width, 1), + args_h.view(1, 1, height, 1, -1).repeat(batch_size, duration, 1, width, 1), + args_w.view(1, 1, 1, width, -1).repeat(batch_size, duration, height, 1, 1), ], dim=-1, ) @@ -258,9 +240,7 @@ class Kandinsky5Modulation(nn.Module): def __init__(self, time_dim: int, model_dim: int, num_params: int): super().__init__() self.activation = nn.SiLU() - self.out_layer = ReplicatedLinear(time_dim, - num_params * model_dim, - bias=True) + self.out_layer = ReplicatedLinear(time_dim, num_params * model_dim, bias=True) self.out_layer.weight.data.zero_() self.out_layer.bias.data.zero_() @@ -395,12 +375,9 @@ class Kandinsky5Attention(nn.Module): query_shape = query.shape[:-2] key_shape = key.shape[:-2] - query = query.reshape(query_shape[0], -1, self.num_heads, - query.shape[-1]).transpose(1, 2) - key = key.reshape(key_shape[0], -1, self.num_heads, - key.shape[-1]).transpose(1, 2) - value = value.reshape(key_shape[0], -1, self.num_heads, - value.shape[-1]).transpose(1, 2) + query = query.reshape(query_shape[0], -1, self.num_heads, query.shape[-1]).transpose(1, 2) + key = key.reshape(key_shape[0], -1, self.num_heads, key.shape[-1]).transpose(1, 2) + value = value.reshape(key_shape[0], -1, self.num_heads, value.shape[-1]).transpose(1, 2) hidden_states = F.scaled_dot_product_attention( query, key, @@ -408,9 +385,8 @@ class Kandinsky5Attention(nn.Module): attn_mask=None, is_causal=False, ) - hidden_states = hidden_states.transpose(1, 2).reshape( - *query_shape, self.num_heads, -1) - + hidden_states = hidden_states.transpose(1, 2).reshape(*query_shape, self.num_heads, -1) + hidden_states = hidden_states.flatten(-2, -1) hidden_states, _ = self.out_layer(hidden_states) @@ -419,11 +395,9 @@ class Kandinsky5Attention(nn.Module): class Kandinsky5FeedForward(nn.Module): - def __init__(self, dim: int, ff_dim: int, prefix: str = "", - quant_config: QuantizationConfig | None = None): + def __init__(self, dim: int, ff_dim: int, prefix: str = "", quant_config: QuantizationConfig | None = None): super().__init__() - self.mlp = MLP(dim, ff_dim, bias=False, act_type="gelu", - quant_config=quant_config, prefix=f"{prefix}.mlp") + self.mlp = MLP(dim, ff_dim, bias=False, act_type="gelu", quant_config=quant_config, prefix=f"{prefix}.mlp") def forward(self, x: torch.Tensor) -> torch.Tensor: return self.mlp(x) @@ -431,26 +405,17 @@ class Kandinsky5FeedForward(nn.Module): class Kandinsky5OutLayer(nn.Module): - def __init__(self, model_dim: int, time_dim: int, visual_dim: int, - patch_size: tuple[int, int, int]): + def __init__(self, model_dim: int, time_dim: int, visual_dim: int, patch_size: tuple[int, int, int]): super().__init__() self.patch_size = patch_size self.modulation = Kandinsky5Modulation(time_dim, model_dim, 2) - self.norm = nn.LayerNorm(model_dim, - eps=1e-5, - elementwise_affine=False) - self.out_layer = ReplicatedLinear(model_dim, - math.prod(patch_size) * visual_dim, - bias=True) + self.norm = nn.LayerNorm(model_dim, eps=1e-5, elementwise_affine=False) + self.out_layer = ReplicatedLinear(model_dim, math.prod(patch_size) * visual_dim, bias=True) - def forward(self, visual_embed: torch.Tensor, - time_embed: torch.Tensor) -> torch.Tensor: - shift, scale = torch.chunk( - self.modulation(time_embed).unsqueeze(dim=1), 2, dim=-1) - visual_embed = ( - self.norm(visual_embed.float()) * - (scale.float()[:, None, None] + 1.0) + - shift.float()[:, None, None]).type_as(visual_embed) + def forward(self, visual_embed: torch.Tensor, time_embed: torch.Tensor) -> torch.Tensor: + shift, scale = torch.chunk(self.modulation(time_embed).unsqueeze(dim=1), 2, dim=-1) + visual_embed = (self.norm(visual_embed.float()) * (scale.float()[:, None, None] + 1.0) + + shift.float()[:, None, None]).type_as(visual_embed) x, _ = self.out_layer(visual_embed) @@ -464,14 +429,16 @@ class Kandinsky5OutLayer(nn.Module): self.patch_size[0], self.patch_size[1], self.patch_size[2], - ).permute(0, 1, 5, 2, 6, 3, 7, 4).flatten(1, 2).flatten( - 2, 3).flatten(3, 4)) + ).permute(0, 1, 5, 2, 6, 3, 7, 4).flatten(1, 2).flatten(2, 3).flatten(3, 4)) return x class Kandinsky5TransformerEncoderBlock(nn.Module): - def __init__(self, model_dim: int, time_dim: int, ff_dim: int, + def __init__(self, + model_dim: int, + time_dim: int, + ff_dim: int, head_dim: int, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, @@ -480,47 +447,38 @@ class Kandinsky5TransformerEncoderBlock(nn.Module): super().__init__() self.text_modulation = Kandinsky5Modulation(time_dim, model_dim, 6) - self.self_attention_norm = LayerNormScaleShift( - model_dim, - norm_type="layer", - eps=1e-5, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) - self.self_attention = Kandinsky5Attention( - model_dim, - head_dim, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.self_attention", - quant_config=quant_config) + self.self_attention_norm = LayerNormScaleShift(model_dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.self_attention = Kandinsky5Attention(model_dim, + head_dim, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.self_attention", + quant_config=quant_config) - self.feed_forward_norm = LayerNormScaleShift( - model_dim, - norm_type="layer", - eps=1e-5, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) - self.feed_forward = Kandinsky5FeedForward( - model_dim, ff_dim, prefix=f"{prefix}.feed_forward", - quant_config=quant_config) + self.feed_forward_norm = LayerNormScaleShift(model_dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.feed_forward = Kandinsky5FeedForward(model_dim, + ff_dim, + prefix=f"{prefix}.feed_forward", + quant_config=quant_config) - def forward(self, x: torch.Tensor, time_embed: torch.Tensor, - rope: torch.Tensor) -> torch.Tensor: - self_attn_params, ff_params = torch.chunk( - self.text_modulation(time_embed).unsqueeze(dim=1), 2, dim=-1) + def forward(self, x: torch.Tensor, time_embed: torch.Tensor, rope: torch.Tensor) -> torch.Tensor: + self_attn_params, ff_params = torch.chunk(self.text_modulation(time_embed).unsqueeze(dim=1), 2, dim=-1) shift, scale, gate = torch.chunk(self_attn_params, 3, dim=-1) - out = self.self_attention_norm(x.float(), - shift=shift, - scale=scale, - convert_modulation_dtype=True).type_as(x) + out = self.self_attention_norm(x.float(), shift=shift, scale=scale, convert_modulation_dtype=True).type_as(x) out = self.self_attention(out, rotary_emb=rope) x = (x.float() + gate.float() * out.float()).type_as(x) ff_shift, ff_scale, ff_gate = torch.chunk(ff_params, 3, dim=-1) - out = self.feed_forward_norm(x.float(), - shift=ff_shift, - scale=ff_scale, + out = self.feed_forward_norm(x.float(), shift=ff_shift, scale=ff_scale, convert_modulation_dtype=True).type_as(x) out = self.feed_forward(out) x = (x.float() + ff_gate.float() * out.float()).type_as(x) @@ -530,7 +488,10 @@ class Kandinsky5TransformerEncoderBlock(nn.Module): class Kandinsky5TransformerDecoderBlock(nn.Module): - def __init__(self, model_dim: int, time_dim: int, ff_dim: int, + def __init__(self, + model_dim: int, + time_dim: int, + ff_dim: int, head_dim: int, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, @@ -540,82 +501,66 @@ class Kandinsky5TransformerDecoderBlock(nn.Module): super().__init__() self.visual_modulation = Kandinsky5Modulation(time_dim, model_dim, 9) - self.self_attention_norm = LayerNormScaleShift( - model_dim, - norm_type="layer", - eps=1e-5, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) - self.self_attention = Kandinsky5Attention( - model_dim, - head_dim, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.self_attention", - use_nabla=use_nabla, - quant_config=quant_config) + self.self_attention_norm = LayerNormScaleShift(model_dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.self_attention = Kandinsky5Attention(model_dim, + head_dim, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.self_attention", + use_nabla=use_nabla, + quant_config=quant_config) - self.cross_attention_norm = LayerNormScaleShift( - model_dim, - norm_type="layer", - eps=1e-5, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) - self.cross_attention = Kandinsky5Attention( - model_dim, - head_dim, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.cross_attention", - quant_config=quant_config) + self.cross_attention_norm = LayerNormScaleShift(model_dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.cross_attention = Kandinsky5Attention(model_dim, + head_dim, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.cross_attention", + quant_config=quant_config) - self.feed_forward_norm = LayerNormScaleShift( - model_dim, - norm_type="layer", - eps=1e-5, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) - self.feed_forward = Kandinsky5FeedForward( - model_dim, ff_dim, prefix=f"{prefix}.feed_forward", - quant_config=quant_config) + self.feed_forward_norm = LayerNormScaleShift(model_dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.feed_forward = Kandinsky5FeedForward(model_dim, + ff_dim, + prefix=f"{prefix}.feed_forward", + quant_config=quant_config) - def forward(self, visual_embed: torch.Tensor, text_embed: torch.Tensor, - time_embed: torch.Tensor, rope: torch.Tensor, - sparse_params: dict[str, Any] | None) -> torch.Tensor: + def forward(self, visual_embed: torch.Tensor, text_embed: torch.Tensor, time_embed: torch.Tensor, + rope: torch.Tensor, sparse_params: dict[str, Any] | None) -> torch.Tensor: self_attn_params, cross_attn_params, ff_params = torch.chunk( self.visual_modulation(time_embed).unsqueeze(dim=1), 3, dim=-1) - self_shift, self_scale, self_gate = torch.chunk(self_attn_params, - 3, - dim=-1) + self_shift, self_scale, self_gate = torch.chunk(self_attn_params, 3, dim=-1) visual_out = self.self_attention_norm( visual_embed.float(), shift=self_shift, scale=self_scale, convert_modulation_dtype=True, ).type_as(visual_embed) - visual_out = self.self_attention(visual_out, - rotary_emb=rope, - sparse_params=sparse_params) - visual_embed = (visual_embed.float() + - self_gate.float() * visual_out.float()).type_as( - visual_embed) + visual_out = self.self_attention(visual_out, rotary_emb=rope, sparse_params=sparse_params) + visual_embed = (visual_embed.float() + self_gate.float() * visual_out.float()).type_as(visual_embed) - cross_shift, cross_scale, cross_gate = torch.chunk(cross_attn_params, - 3, - dim=-1) + cross_shift, cross_scale, cross_gate = torch.chunk(cross_attn_params, 3, dim=-1) visual_out = self.cross_attention_norm( visual_embed.float(), shift=cross_shift, scale=cross_scale, convert_modulation_dtype=True, ).type_as(visual_embed) - visual_out = self.cross_attention( - visual_out, encoder_hidden_states=text_embed) - visual_embed = (visual_embed.float() + - cross_gate.float() * visual_out.float()).type_as( - visual_embed) + visual_out = self.cross_attention(visual_out, encoder_hidden_states=text_embed) + visual_embed = (visual_embed.float() + cross_gate.float() * visual_out.float()).type_as(visual_embed) ff_shift, ff_scale, ff_gate = torch.chunk(ff_params, 3, dim=-1) visual_out = self.feed_forward_norm( @@ -625,9 +570,7 @@ class Kandinsky5TransformerDecoderBlock(nn.Module): convert_modulation_dtype=True, ).type_as(visual_embed) visual_out = self.feed_forward(visual_out) - visual_embed = (visual_embed.float() + - ff_gate.float() * visual_out.float()).type_as( - visual_embed) + visual_embed = (visual_embed.float() + ff_gate.float() * visual_out.float()).type_as(visual_embed) return visual_embed @@ -649,8 +592,7 @@ class Kandinsky5Transformer3DModel(BaseDiT): lora_param_names_mapping = _ARCH_CONFIG_DEFAULTS.lora_param_names_mapping _supported_attention_backends = _ARCH_CONFIG_DEFAULTS._supported_attention_backends - def __init__(self, config: Kandinsky5VideoConfig, - hf_config: dict[str, Any]) -> None: + def __init__(self, config: Kandinsky5VideoConfig, hf_config: dict[str, Any]) -> None: super().__init__(config=config, hf_config=hf_config) arch = config.arch_config quant_config = config.quant_config @@ -663,44 +605,37 @@ class Kandinsky5Transformer3DModel(BaseDiT): self.visual_cond = arch.visual_cond self.attention_type = arch.attention_type - visual_embed_dim = (2 * arch.in_visual_dim + - 1) if arch.visual_cond else arch.in_visual_dim + visual_embed_dim = (2 * arch.in_visual_dim + 1) if arch.visual_cond else arch.in_visual_dim - self.time_embeddings = Kandinsky5TimeEmbeddings( - arch.model_dim, arch.time_dim) - self.text_embeddings = Kandinsky5TextEmbeddings( - arch.in_text_dim, arch.model_dim) - self.pooled_text_embeddings = Kandinsky5TextEmbeddings( - arch.in_text_dim2, arch.time_dim) - self.visual_embeddings = Kandinsky5VisualEmbeddings( - visual_embed_dim, arch.model_dim, arch.patch_size) + self.time_embeddings = Kandinsky5TimeEmbeddings(arch.model_dim, arch.time_dim) + self.text_embeddings = Kandinsky5TextEmbeddings(arch.in_text_dim, arch.model_dim) + self.pooled_text_embeddings = Kandinsky5TextEmbeddings(arch.in_text_dim2, arch.time_dim) + self.visual_embeddings = Kandinsky5VisualEmbeddings(visual_embed_dim, arch.model_dim, arch.patch_size) self.text_rope_embeddings = Kandinsky5RoPE1D(head_dim) self.visual_rope_embeddings = Kandinsky5RoPE3D(arch.axes_dims) self.text_transformer_blocks = nn.ModuleList([ - Kandinsky5TransformerEncoderBlock(arch.model_dim, arch.time_dim, + Kandinsky5TransformerEncoderBlock(arch.model_dim, + arch.time_dim, arch.ff_dim, head_dim, self._supported_attention_backends, prefix=f"{config.prefix}.text_transformer_blocks.{i}", - quant_config=quant_config) - for i in range(arch.num_text_blocks) + quant_config=quant_config) for i in range(arch.num_text_blocks) ]) self.visual_transformer_blocks = nn.ModuleList([ - Kandinsky5TransformerDecoderBlock(arch.model_dim, arch.time_dim, + Kandinsky5TransformerDecoderBlock(arch.model_dim, + arch.time_dim, arch.ff_dim, head_dim, self._supported_attention_backends, prefix=f"{config.prefix}.visual_transformer_blocks.{i}", use_nabla=arch.attention_type == "nabla", - quant_config=quant_config) - for i in range(arch.num_visual_blocks) + quant_config=quant_config) for i in range(arch.num_visual_blocks) ]) - self.out_layer = Kandinsky5OutLayer(arch.model_dim, arch.time_dim, - arch.out_visual_dim, - arch.patch_size) + self.out_layer = Kandinsky5OutLayer(arch.model_dim, arch.time_dim, arch.out_visual_dim, arch.patch_size) self.gradient_checkpointing = False self.hidden_size = arch.hidden_size @@ -716,8 +651,7 @@ class Kandinsky5Transformer3DModel(BaseDiT): encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, pooled_projections: torch.Tensor | None = None, - visual_rope_pos: tuple[torch.Tensor, torch.Tensor, - torch.Tensor] | list[torch.Tensor] | None = None, + visual_rope_pos: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | list[torch.Tensor] | None = None, text_rope_pos: torch.Tensor | None = None, scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), sparse_params: dict[str, Any] | None = None, @@ -729,16 +663,13 @@ class Kandinsky5Transformer3DModel(BaseDiT): if pooled_projections is None: if encoder_hidden_states_image is None: - raise ValueError( - "pooled_projections must be provided for Kandinsky5.") + raise ValueError("pooled_projections must be provided for Kandinsky5.") pooled_projections = encoder_hidden_states_image if isinstance(pooled_projections, list): pooled_projections = pooled_projections[0] if visual_rope_pos is None or text_rope_pos is None: - raise ValueError( - "visual_rope_pos and text_rope_pos are required for Kandinsky5." - ) + raise ValueError("visual_rope_pos and text_rope_pos are required for Kandinsky5.") x = hidden_states text_embed = self.text_embeddings(encoder_hidden_states) @@ -749,36 +680,29 @@ class Kandinsky5Transformer3DModel(BaseDiT): for text_transformer_block in self.text_transformer_blocks: if torch.is_grad_enabled() and self.gradient_checkpointing: - text_embed = torch.utils.checkpoint.checkpoint( - text_transformer_block, - text_embed, - time_embed, - text_rope, - use_reentrant=False) + text_embed = torch.utils.checkpoint.checkpoint(text_transformer_block, + text_embed, + time_embed, + text_rope, + use_reentrant=False) else: - text_embed = text_transformer_block(text_embed, time_embed, - text_rope) + text_embed = text_transformer_block(text_embed, time_embed, text_rope) visual_shape = visual_embed.shape[:-1] - visual_rope = self.visual_rope_embeddings(visual_shape, visual_rope_pos, - scale_factor) - to_fractal = sparse_params[ - "to_fractal"] if sparse_params is not None else False - - visual_embed, visual_rope = fractal_flatten(visual_embed, visual_rope, - visual_shape, - block_mask=to_fractal) + visual_rope = self.visual_rope_embeddings(visual_shape, visual_rope_pos, scale_factor) + to_fractal = sparse_params["to_fractal"] if sparse_params is not None else False + + visual_embed, visual_rope = fractal_flatten(visual_embed, visual_rope, visual_shape, block_mask=to_fractal) for visual_transformer_block in self.visual_transformer_blocks: if torch.is_grad_enabled() and self.gradient_checkpointing: - visual_embed = torch.utils.checkpoint.checkpoint( - visual_transformer_block, - visual_embed, - text_embed, - time_embed, - visual_rope, - sparse_params, - use_reentrant=False) + visual_embed = torch.utils.checkpoint.checkpoint(visual_transformer_block, + visual_embed, + text_embed, + time_embed, + visual_rope, + sparse_params, + use_reentrant=False) else: visual_embed = visual_transformer_block( visual_embed, @@ -788,46 +712,30 @@ class Kandinsky5Transformer3DModel(BaseDiT): sparse_params, ) - visual_embed = fractal_unflatten(visual_embed, - visual_shape, - block_mask=to_fractal) + visual_embed = fractal_unflatten(visual_embed, visual_shape, block_mask=to_fractal) x = self.out_layer(visual_embed, time_embed) if return_dict: return Kandinsky5TransformerOutput(sample=x) - + return x - def materialize_non_persistent_buffers(self, device: torch.device, - dtype: torch.dtype | None = None - ) -> None: - if isinstance(self.time_embeddings.freqs, - torch.Tensor) and self.time_embeddings.freqs.is_meta: - self.time_embeddings.freqs = _build_rotary_freqs( - self.time_embeddings.model_dim // 2, - self.time_embeddings.max_period).to(device=device) + def materialize_non_persistent_buffers(self, device: torch.device, dtype: torch.dtype | None = None) -> None: + if isinstance(self.time_embeddings.freqs, torch.Tensor) and self.time_embeddings.freqs.is_meta: + self.time_embeddings.freqs = _build_rotary_freqs(self.time_embeddings.model_dim // 2, + self.time_embeddings.max_period).to(device=device) - if isinstance(self.text_rope_embeddings.args, - torch.Tensor) and self.text_rope_embeddings.args.is_meta: + if isinstance(self.text_rope_embeddings.args, torch.Tensor) and self.text_rope_embeddings.args.is_meta: freq = _build_rotary_freqs(self.text_rope_embeddings.dim // 2, - self.text_rope_embeddings.max_period).to( - device=device) - pos = torch.arange(self.text_rope_embeddings.max_pos, - dtype=freq.dtype, - device=device) + self.text_rope_embeddings.max_period).to(device=device) + pos = torch.arange(self.text_rope_embeddings.max_pos, dtype=freq.dtype, device=device) self.text_rope_embeddings._buffers["args"] = torch.outer(pos, freq) - for i, (axes_dim, - ax_max_pos) in enumerate( - zip(self.visual_rope_embeddings.axes_dims, - self.visual_rope_embeddings.max_pos, - strict=True)): + for i, (axes_dim, ax_max_pos) in enumerate( + zip(self.visual_rope_embeddings.axes_dims, self.visual_rope_embeddings.max_pos, strict=True)): name = f"args_{i}" buf = getattr(self.visual_rope_embeddings, name, None) if isinstance(buf, torch.Tensor) and buf.is_meta: - freq = _build_rotary_freqs( - axes_dim // 2, self.visual_rope_embeddings.max_period).to( - device=device) + freq = _build_rotary_freqs(axes_dim // 2, self.visual_rope_embeddings.max_period).to(device=device) pos = torch.arange(ax_max_pos, dtype=freq.dtype, device=device) - self.visual_rope_embeddings._buffers[name] = torch.outer( - pos, freq) + self.visual_rope_embeddings._buffers[name] = torch.outer(pos, freq) diff --git a/fastvideo/models/dits/lingbotworld/cam_utils.py b/fastvideo/models/dits/lingbotworld/cam_utils.py index 3ac497b6e..9190ce27e 100644 --- a/fastvideo/models/dits/lingbotworld/cam_utils.py +++ b/fastvideo/models/dits/lingbotworld/cam_utils.py @@ -7,21 +7,21 @@ import torch from scipy.interpolate import interp1d from scipy.spatial.transform import Rotation, Slerp - # --- Official Code (Leave Unchanged) --- - + + def interpolate_camera_poses( - src_indices: np.ndarray, - src_rot_mat: np.ndarray, - src_trans_vec: np.ndarray, + src_indices: np.ndarray, + src_rot_mat: np.ndarray, + src_trans_vec: np.ndarray, tgt_indices: np.ndarray, ) -> torch.Tensor: # interpolate translation interp_func_trans = interp1d( - src_indices, - src_trans_vec, - axis=0, - kind='linear', + src_indices, + src_trans_vec, + axis=0, + kind='linear', bounds_error=False, fill_value="extrapolate", ) @@ -32,7 +32,7 @@ def interpolate_camera_poses( # ensure there is no sudden change in qw quats = src_quat_vec.as_quat().copy() # [N, 4] for i in range(1, len(quats)): - if np.dot(quats[i], quats[i-1]) < 0: + if np.dot(quats[i], quats[i - 1]) < 0: quats[i] = -quats[i] src_quat_vec = Rotation.from_quat(quats) slerp_func_rot = Slerp(src_indices, src_quat_vec) @@ -47,8 +47,8 @@ def interpolate_camera_poses( def SE3_inverse(T: torch.Tensor) -> torch.Tensor: - Rot = T[:, :3, :3] # [B,3,3] - trans = T[:, :3, 3:] # [B,3,1] + Rot = T[:, :3, :3] # [B,3,3] + trans = T[:, :3, 3:] # [B,3,1] R_inv = Rot.transpose(-1, -2) t_inv = -torch.bmm(R_inv, trans) T_inv = torch.eye(4, device=T.device, dtype=T.dtype)[None, :, :].repeat(T.shape[0], 1, 1) @@ -58,9 +58,9 @@ def SE3_inverse(T: torch.Tensor) -> torch.Tensor: def compute_relative_poses( - c2ws_mat: torch.Tensor, - framewise: bool = False, - normalize_trans: bool = True, + c2ws_mat: torch.Tensor, + framewise: bool = False, + normalize_trans: bool = True, ) -> torch.Tensor: ref_w2cs = SE3_inverse(c2ws_mat[0:1]) relative_poses = torch.matmul(ref_w2cs, c2ws_mat) @@ -70,8 +70,8 @@ def compute_relative_poses( # compute pose between i and i+1 relative_poses_framewise = torch.bmm(SE3_inverse(relative_poses[:-1]), relative_poses[1:]) relative_poses[1:] = relative_poses_framewise - if normalize_trans: # note refer to camctrl2: "we scale the coordinate inputs to roughly 1 standard deviation to simplify model learning." - translations = relative_poses[:, :3, 3] # [f, 3] + if normalize_trans: # note refer to camctrl2: "we scale the coordinate inputs to roughly 1 standard deviation to simplify model learning." + translations = relative_poses[:, :3, 3] # [f, 3] max_norm = torch.norm(translations, dim=-1).max() # only normlaize when moving if max_norm > 0: @@ -80,12 +80,17 @@ def compute_relative_poses( @torch.no_grad() -def create_meshgrid(n_frames: int, height: int, width: int, bias: float = 0.5, device='cuda', dtype=torch.float32) -> torch.Tensor: +def create_meshgrid(n_frames: int, + height: int, + width: int, + bias: float = 0.5, + device='cuda', + dtype=torch.float32) -> torch.Tensor: x_range = torch.arange(width, device=device, dtype=dtype) y_range = torch.arange(height, device=device, dtype=dtype) grid_y, grid_x = torch.meshgrid(y_range, x_range, indexing='ij') - grid_xy = torch.stack([grid_x, grid_y], dim=-1).view([-1, 2]) + bias # [h*w, 2] - grid_xy = grid_xy[None, ...].repeat(n_frames, 1, 1) # [f, h*w, 2] + grid_xy = torch.stack([grid_x, grid_y], dim=-1).view([-1, 2]) + bias # [h*w, 2] + grid_xy = grid_xy[None, ...].repeat(n_frames, 1, 1) # [f, h*w, 2] return grid_xy @@ -96,25 +101,25 @@ def get_plucker_embeddings( width: int, ): n_frames = c2ws_mat.shape[0] - grid_xy = create_meshgrid(n_frames, height, width, device=c2ws_mat.device, dtype=c2ws_mat.dtype) # [f, h*w, 2] - fx, fy, cx, cy = Ks.chunk(4, dim=-1) # [f, 1] + grid_xy = create_meshgrid(n_frames, height, width, device=c2ws_mat.device, dtype=c2ws_mat.dtype) # [f, h*w, 2] + fx, fy, cx, cy = Ks.chunk(4, dim=-1) # [f, 1] - i = grid_xy[..., 0] # [f, h*w] - j = grid_xy[..., 1] # [f, h*w] - zs = torch.ones_like(i) # [f, h*w] + i = grid_xy[..., 0] # [f, h*w] + j = grid_xy[..., 1] # [f, h*w] + zs = torch.ones_like(i) # [f, h*w] xs = (i - cx) / fx * zs ys = (j - cy) / fy * zs - directions = torch.stack([xs, ys, zs], dim=-1) # [f, h*w, 3] - directions = directions / directions.norm(dim=-1, keepdim=True) # [f, h*w, 3] + directions = torch.stack([xs, ys, zs], dim=-1) # [f, h*w, 3] + directions = directions / directions.norm(dim=-1, keepdim=True) # [f, h*w, 3] - rays_d = directions @ c2ws_mat[:, :3, :3].transpose(-1, -2) # [f, h*w, 3] - rays_o = c2ws_mat[:, :3, 3] # [f, 3] - rays_o = rays_o[:, None, :].expand_as(rays_d) # [f, h*w, 3] + rays_d = directions @ c2ws_mat[:, :3, :3].transpose(-1, -2) # [f, h*w, 3] + rays_o = c2ws_mat[:, :3, 3] # [f, 3] + rays_o = rays_o[:, None, :].expand_as(rays_d) # [f, h*w, 3] # rays_dxo = torch.cross(rays_o, rays_d, dim=-1) # [f, h*w, 3] # note refer to: apt2 - plucker_embeddings = torch.cat([rays_o, rays_d], dim=-1) # [f, h*w, 6] - plucker_embeddings = plucker_embeddings.view([n_frames, height, width, 6]) # [f*h*w, 6] + plucker_embeddings = torch.cat([rays_o, rays_d], dim=-1) # [f, h*w, 6] + plucker_embeddings = plucker_embeddings.view([n_frames, height, width, 6]) # [f*h*w, 6] return plucker_embeddings @@ -127,7 +132,7 @@ def get_Ks_transformed( height_final: int, width_final: int, ): - fx, fy, cx, cy = Ks.chunk(4, dim=-1) # [f, 1] + fx, fy, cx, cy = Ks.chunk(4, dim=-1) # [f, 1] scale_x = width_resize / width_org scale_y = height_resize / height_org @@ -142,7 +147,7 @@ def get_Ks_transformed( cx_final = cx_resize - crop_offset_x cy_final = cy_resize - crop_offset_y - + Ks_transformed = torch.zeros_like(Ks) Ks_transformed[:, 0:1] = fx_resize Ks_transformed[:, 1:2] = fy_resize @@ -154,6 +159,7 @@ def get_Ks_transformed( # --- Custom --- + def prepare_camera_embedding( action_path: str, num_frames: int, @@ -166,9 +172,7 @@ def prepare_camera_embedding( num_frames = min(num_frames, len_c2ws) c2ws = c2ws[:num_frames] - Ks = torch.from_numpy( - np.load(os.path.join(action_path, "intrinsics.npy")) - ).float() + Ks = torch.from_numpy(np.load(os.path.join(action_path, "intrinsics.npy"))).float() Ks = get_Ks_transformed( Ks, height_org=480, diff --git a/fastvideo/models/dits/lingbotworld/model.py b/fastvideo/models/dits/lingbotworld/model.py index 56c51e972..26009396a 100644 --- a/fastvideo/models/dits/lingbotworld/model.py +++ b/fastvideo/models/dits/lingbotworld/model.py @@ -8,11 +8,9 @@ import torch.nn as nn from fastvideo.attention import DistributedAttention from fastvideo.configs.models.dits.lingbotworld import LingBotWorldVideoConfig -from fastvideo.distributed.communication_op import ( - sequence_model_parallel_all_gather_with_unpad, - sequence_model_parallel_shard) -from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, - RMSNorm, ScaleResidual, +from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad, + sequence_model_parallel_shard) +from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual, ScaleResidualLayerNormScaleShift) from fastvideo.layers.linear import ReplicatedLinear # from torch.nn import RMSNorm @@ -51,8 +49,7 @@ class LingBotWorldCamConditioner(nn.Module): return hidden_states assert c2ws_plucker_emb.shape == hidden_states.shape, ( f"c2ws_plucker_emb shape must match hidden_states shape, got " - f"{tuple(c2ws_plucker_emb.shape)} vs {tuple(hidden_states.shape)}" - ) + f"{tuple(c2ws_plucker_emb.shape)} vs {tuple(hidden_states.shape)}") c2ws_hidden_states = self.cam_injector(c2ws_plucker_emb) c2ws_hidden_states = c2ws_hidden_states + c2ws_plucker_emb cam_scale = self.cam_scale_layer(c2ws_hidden_states) @@ -82,12 +79,11 @@ class LingBotWorldTransformerBlock(nn.Module): self.to_v = ReplicatedLinear(dim, dim, bias=True) self.to_out = ReplicatedLinear(dim, dim, bias=True) - self.attn1 = DistributedAttention( - num_heads=num_heads, - head_size=dim // num_heads, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.attn1") + self.attn1 = DistributedAttention(num_heads=num_heads, + head_size=dim // num_heads, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn1") self.hidden_dim = dim self.num_attention_heads = num_heads dim_head = dim // num_heads @@ -99,37 +95,28 @@ class LingBotWorldTransformerBlock(nn.Module): self.norm_q = RMSNorm(dim, eps=eps) self.norm_k = RMSNorm(dim, eps=eps) else: - raise NotImplementedError( - f"QK Norm type '{qk_norm}' not supported") + raise NotImplementedError(f"QK Norm type '{qk_norm}' not supported") assert cross_attn_norm is True - self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( - dim, - norm_type="layer", - eps=eps, - elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32) + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32, + compute_dtype=torch.float32) # 2. Cross-attention if added_kv_proj_dim is not None: # I2V - self.attn2 = WanI2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps) + self.attn2 = WanI2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) else: # T2V - self.attn2 = WanT2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps) - self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( - dim, - norm_type="layer", - eps=eps, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") @@ -156,8 +143,7 @@ class LingBotWorldTransformerBlock(nn.Module): if temb.dim() == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - self.scale_shift_table.unsqueeze(0) + temb.float() - ).chunk(6, dim=2) + self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2) # batch_size, seq_len, 1, inner_dim shift_msa = shift_msa.squeeze(2) scale_msa = scale_msa.squeeze(2) @@ -168,13 +154,11 @@ class LingBotWorldTransformerBlock(nn.Module): else: # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) e = self.scale_shift_table + temb.float() - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( - 6, dim=1) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1) assert shift_msa.dtype == torch.float32 # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()) * - (1 + scale_msa) + shift_msa).to(orig_dtype) + norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -200,10 +184,9 @@ class LingBotWorldTransformerBlock(nn.Module): attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.tensor([0], device=hidden_states.device) - norm_hidden_states, hidden_states = self.self_attn_residual_norm( - hidden_states, attn_output, gate_msa, null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, + null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) # Inject camera condition # must be applied after the self-attention residual update. hidden_states = self.cam_conditioner(hidden_states, c2ws_plucker_emb) @@ -211,13 +194,10 @@ class LingBotWorldTransformerBlock(nn.Module): norm_hidden_states = norm_hidden_states.to(orig_dtype) # 2. Cross-attention - attn_output = self.attn2(norm_hidden_states, - context=encoder_hidden_states, - context_lens=None) - norm_hidden_states, hidden_states = self.cross_attn_residual_norm( - hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) + attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None) + norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa, + c_scale_msa) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) @@ -230,14 +210,12 @@ class LingBotWorldTransformerBlock(nn.Module): class LingBotWorldTransformer3DModel(BaseDiT): _fsdp_shard_conditions = LingBotWorldVideoConfig()._fsdp_shard_conditions _compile_conditions = LingBotWorldVideoConfig()._compile_conditions - _supported_attention_backends = LingBotWorldVideoConfig( - )._supported_attention_backends + _supported_attention_backends = LingBotWorldVideoConfig()._supported_attention_backends param_names_mapping = LingBotWorldVideoConfig().param_names_mapping reverse_param_names_mapping = LingBotWorldVideoConfig().reverse_param_names_mapping lora_param_names_mapping = LingBotWorldVideoConfig().lora_param_names_mapping - def __init__(self, config: LingBotWorldVideoConfig, hf_config: dict[str, - Any]) -> None: + def __init__(self, config: LingBotWorldVideoConfig, hf_config: dict[str, Any]) -> None: super().__init__(config=config, hf_config=hf_config) inner_dim = config.num_attention_heads * config.attention_head_dim @@ -249,7 +227,8 @@ class LingBotWorldTransformer3DModel(BaseDiT): self.patch_size = config.patch_size self.text_len = config.text_len - assert config.num_attention_heads % get_sp_world_size() == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})" + assert config.num_attention_heads % get_sp_world_size( + ) == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})" # 1. Patch & position embedding self.patch_embedding = PatchEmbed(in_chans=config.in_channels, @@ -280,8 +259,7 @@ class LingBotWorldTransformer3DModel(BaseDiT): config.eps, config.added_kv_proj_dim, self._supported_attention_backends, - prefix=f"{config.prefix}.blocks.{i}") - for i in range(config.num_layers) + prefix=f"{config.prefix}.blocks.{i}") for i in range(config.num_layers) ]) # 4. Output norm & projection @@ -291,10 +269,8 @@ class LingBotWorldTransformer3DModel(BaseDiT): elementwise_affine=False, dtype=torch.float32, compute_dtype=torch.float32) - self.proj_out = nn.Linear( - inner_dim, config.out_channels * math.prod(config.patch_size)) - self.scale_shift_table = nn.Parameter( - torch.randn(1, 2, inner_dim) / inner_dim**0.5) + self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) self.gradient_checkpointing = False self.__post_init__() @@ -324,34 +300,29 @@ class LingBotWorldTransformer3DModel(BaseDiT): # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] - freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames, post_patch_height, - post_patch_width), - self.hidden_size, - self.num_attention_heads, - rope_dim_list, - dtype=torch.float32 if current_platform.is_mps() else torch.float64, - rope_theta=10000) - freqs_cis = (freqs_cos.to(hidden_states.device).float(), - freqs_sin.to(hidden_states.device).float()) + freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000) + freqs_cis = (freqs_cos.to(hidden_states.device).float(), freqs_sin.to(hidden_states.device).float()) hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) c2ws_hidden_states = None if c2ws_plucker_emb is not None: c2ws_plucker_emb = self.patch_embedding_wancamctrl( - c2ws_plucker_emb.to(device=hidden_states.device, dtype=hidden_states.dtype) - ) + c2ws_plucker_emb.to(device=hidden_states.device, dtype=hidden_states.dtype)) c2ws_hidden_states = self.c2ws_mlp(c2ws_plucker_emb) c2ws_plucker_emb = c2ws_plucker_emb + c2ws_hidden_states # Shard with padding support - returns (sharded_tensor, original_seq_len) hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1) - + # Shard c2ws_plucker_emb if c2ws_plucker_emb is not None: c2ws_plucker_emb, _ = sequence_model_parallel_shard(c2ws_plucker_emb, dim=1) - # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v) if timestep.dim() == 2: @@ -371,28 +342,26 @@ class LingBotWorldTransformer3DModel(BaseDiT): if encoder_hidden_states_image is not None: if encoder_hidden_states is not None: - encoder_hidden_states = torch.concat( - [encoder_hidden_states_image, encoder_hidden_states], dim=1) + encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) else: encoder_hidden_states = encoder_hidden_states_image if current_platform.is_mps() or current_platform.is_npu(): encoder_hidden_states = encoder_hidden_states.to(orig_dtype) else: - encoder_hidden_states = encoder_hidden_states # cast to orig_dtype for MPS & NPU + encoder_hidden_states = encoder_hidden_states # cast to orig_dtype for MPS & NPU assert encoder_hidden_states.dtype == orig_dtype # 4. Transformer blocks if torch.is_grad_enabled() and self.gradient_checkpointing: for block in self.blocks: - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, original_seq_len, c2ws_plucker_emb) + hidden_states = self._gradient_checkpointing_func(block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, original_seq_len, + c2ws_plucker_emb) else: for block in self.blocks: - hidden_states = block(hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, original_seq_len, + hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, original_seq_len, c2ws_plucker_emb) # 5. Output norm, projection & unpatchify if temb.dim() == 3: @@ -403,23 +372,20 @@ class LingBotWorldTransformer3DModel(BaseDiT): else: # batch_size, inner_dim shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) - hidden_states = self.norm_out(hidden_states, shift, scale) # Gather and unpad in one operation - hidden_states = sequence_model_parallel_all_gather_with_unpad( - hidden_states, original_seq_len, dim=1) + hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, - post_patch_height, - post_patch_width, p_t, p_h, p_w, - -1) + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width, + p_t, p_h, p_w, -1) hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) return output + # Entry point for model registry EntryClass = LingBotWorldTransformer3DModel diff --git a/fastvideo/models/dits/lingbotworld2/cam_utils.py b/fastvideo/models/dits/lingbotworld2/cam_utils.py index 30521ddaa..50cb3d018 100644 --- a/fastvideo/models/dits/lingbotworld2/cam_utils.py +++ b/fastvideo/models/dits/lingbotworld2/cam_utils.py @@ -5,21 +5,21 @@ import torch from scipy.interpolate import interp1d from scipy.spatial.transform import Rotation, Slerp - # --- Official Code (Leave Unchanged) --- - + + def interpolate_camera_poses( - src_indices: np.ndarray, - src_rot_mat: np.ndarray, - src_trans_vec: np.ndarray, + src_indices: np.ndarray, + src_rot_mat: np.ndarray, + src_trans_vec: np.ndarray, tgt_indices: np.ndarray, ) -> torch.Tensor: # interpolate translation interp_func_trans = interp1d( - src_indices, - src_trans_vec, - axis=0, - kind='linear', + src_indices, + src_trans_vec, + axis=0, + kind='linear', bounds_error=False, fill_value="extrapolate", ) @@ -30,7 +30,7 @@ def interpolate_camera_poses( # ensure there is no sudden change in qw quats = src_quat_vec.as_quat().copy() # [N, 4] for i in range(1, len(quats)): - if np.dot(quats[i], quats[i-1]) < 0: + if np.dot(quats[i], quats[i - 1]) < 0: quats[i] = -quats[i] src_quat_vec = Rotation.from_quat(quats) slerp_func_rot = Slerp(src_indices, src_quat_vec) @@ -45,8 +45,8 @@ def interpolate_camera_poses( def SE3_inverse(T: torch.Tensor) -> torch.Tensor: - Rot = T[:, :3, :3] # [B,3,3] - trans = T[:, :3, 3:] # [B,3,1] + Rot = T[:, :3, :3] # [B,3,3] + trans = T[:, :3, 3:] # [B,3,1] R_inv = Rot.transpose(-1, -2) t_inv = -torch.bmm(R_inv, trans) T_inv = torch.eye(4, device=T.device, dtype=T.dtype)[None, :, :].repeat(T.shape[0], 1, 1) @@ -56,9 +56,9 @@ def SE3_inverse(T: torch.Tensor) -> torch.Tensor: def compute_relative_poses( - c2ws_mat: torch.Tensor, - framewise: bool = False, - normalize_trans: bool = True, + c2ws_mat: torch.Tensor, + framewise: bool = False, + normalize_trans: bool = True, ) -> torch.Tensor: ref_w2cs = SE3_inverse(c2ws_mat[0:1]) relative_poses = torch.matmul(ref_w2cs, c2ws_mat) @@ -68,8 +68,8 @@ def compute_relative_poses( # compute pose between i and i+1 relative_poses_framewise = torch.bmm(SE3_inverse(relative_poses[:-1]), relative_poses[1:]) relative_poses[1:] = relative_poses_framewise - if normalize_trans: # note refer to camctrl2: "we scale the coordinate inputs to roughly 1 standard deviation to simplify model learning." - translations = relative_poses[:, :3, 3] # [f, 3] + if normalize_trans: # note refer to camctrl2: "we scale the coordinate inputs to roughly 1 standard deviation to simplify model learning." + translations = relative_poses[:, :3, 3] # [f, 3] max_norm = torch.norm(translations, dim=-1).max() # only normlaize when moving if max_norm > 0: @@ -78,12 +78,17 @@ def compute_relative_poses( @torch.no_grad() -def create_meshgrid(n_frames: int, height: int, width: int, bias: float = 0.5, device='cuda', dtype=torch.float32) -> torch.Tensor: +def create_meshgrid(n_frames: int, + height: int, + width: int, + bias: float = 0.5, + device='cuda', + dtype=torch.float32) -> torch.Tensor: x_range = torch.arange(width, device=device, dtype=dtype) y_range = torch.arange(height, device=device, dtype=dtype) grid_y, grid_x = torch.meshgrid(y_range, x_range, indexing='ij') - grid_xy = torch.stack([grid_x, grid_y], dim=-1).view([-1, 2]) + bias # [h*w, 2] - grid_xy = grid_xy[None, ...].repeat(n_frames, 1, 1) # [f, h*w, 2] + grid_xy = torch.stack([grid_x, grid_y], dim=-1).view([-1, 2]) + bias # [h*w, 2] + grid_xy = grid_xy[None, ...].repeat(n_frames, 1, 1) # [f, h*w, 2] return grid_xy @@ -128,7 +133,7 @@ def get_Ks_transformed( height_final: int, width_final: int, ): - fx, fy, cx, cy = Ks.chunk(4, dim=-1) # [f, 1] + fx, fy, cx, cy = Ks.chunk(4, dim=-1) # [f, 1] scale_x = width_resize / width_org scale_y = height_resize / height_org @@ -143,7 +148,7 @@ def get_Ks_transformed( cx_final = cx_resize - crop_offset_x cy_final = cy_resize - crop_offset_y - + Ks_transformed = torch.zeros_like(Ks) Ks_transformed[:, 0:1] = fx_resize Ks_transformed[:, 1:2] = fy_resize @@ -155,6 +160,7 @@ def get_Ks_transformed( # --- Custom --- + def prepare_camera_embedding( action_path: str, num_frames: int, @@ -167,9 +173,7 @@ def prepare_camera_embedding( num_frames = min(num_frames, len_c2ws) c2ws = c2ws[:num_frames] - Ks = torch.from_numpy( - np.load(os.path.join(action_path, "intrinsics.npy")) - ).float() + Ks = torch.from_numpy(np.load(os.path.join(action_path, "intrinsics.npy"))).float() Ks = get_Ks_transformed( Ks, height_org=480, diff --git a/fastvideo/models/dits/lingbotworld2/causal_fast.py b/fastvideo/models/dits/lingbotworld2/causal_fast.py index c87b99f93..e3bac555b 100644 --- a/fastvideo/models/dits/lingbotworld2/causal_fast.py +++ b/fastvideo/models/dits/lingbotworld2/causal_fast.py @@ -11,8 +11,7 @@ import torch.nn as nn import torch.nn.functional as F from fastvideo.configs.models.dits.lingbotworld2 import ( - LingBotWorld2CausalFastVideoConfig, -) + LingBotWorld2CausalFastVideoConfig, ) from fastvideo.distributed.communication_op import ( sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D, @@ -58,7 +57,8 @@ def rope_params(max_seq_len: int, dim: int, theta: int = 10000) -> torch.Tensor: assert dim % 2 == 0 freqs = torch.outer( torch.arange(max_seq_len), - 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)), + 1.0 / torch.pow(theta, + torch.arange(0, dim, 2).to(torch.float64).div(dim)), ) return torch.polar(torch.ones_like(freqs), freqs) @@ -110,12 +110,8 @@ def flash_attention( if version == 3 and not FLASH_ATTN_3_AVAILABLE: warnings.warn("FlashAttention 3 is not available; using FlashAttention 2.") - cu_q = torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum( - 0, dtype=torch.int32 - ).to(q.device, non_blocking=True) - cu_k = torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum( - 0, dtype=torch.int32 - ).to(k.device, non_blocking=True) + cu_q = torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(0, dtype=torch.int32).to(q.device, non_blocking=True) + cu_k = torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(0, dtype=torch.int32).to(k.device, non_blocking=True) if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: x = flash_attn_interface.flash_attn_varlen_func( @@ -189,9 +185,7 @@ def attention( q = q.transpose(1, 2).to(dtype) k = k.transpose(1, 2).to(dtype) v = v.transpose(1, 2).to(dtype) - out = F.scaled_dot_product_attention( - q, k, v, attn_mask=None, is_causal=causal, dropout_p=dropout_p - ) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=causal, dropout_p=dropout_p) return out.transpose(1, 2).contiguous() @@ -211,7 +205,7 @@ def causal_rope_apply( x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)) freqs_i = torch.cat( [ - freqs[0][start_frame : start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), + freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), ], @@ -330,18 +324,19 @@ class CausalWanSelfAttention(nn.Module): local_start_index = current_start kv_cache["k"][:, local_start_index:local_end_index] = roped_key kv_cache["v"][:, local_start_index:local_end_index] = v - elif (current_end > kv_cache["global_end_index"].item()) and ( - num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size - ): + elif (current_end + > kv_cache["global_end_index"].item()) and (num_new_tokens + kv_cache["local_end_index"].item() + > kv_cache_size): num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens - kv_cache["k"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["k"][ - :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens - ].clone() - kv_cache["v"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["v"][ - :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens - ].clone() - local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() - num_evicted_tokens + kv_cache["k"][:, sink_tokens:sink_tokens + + num_rolled_tokens] = kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + + num_evicted_tokens + num_rolled_tokens].clone() + kv_cache["v"][:, sink_tokens:sink_tokens + + num_rolled_tokens] = kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + + num_evicted_tokens + num_rolled_tokens].clone() + local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item( + ) - num_evicted_tokens local_start_index = local_end_index - num_new_tokens kv_cache["k"][:, local_start_index:local_end_index] = roped_key kv_cache["v"][:, local_start_index:local_end_index] = v @@ -351,8 +346,8 @@ class CausalWanSelfAttention(nn.Module): kv_cache["k"][:, local_start_index:local_end_index] = roped_key kv_cache["v"][:, local_start_index:local_end_index] = v - k_cache = kv_cache["k"][:, max(0, local_end_index - max_attention_size) : local_end_index] - v_cache = kv_cache["v"][:, max(0, local_end_index - max_attention_size) : local_end_index] + k_cache = kv_cache["k"][:, max(0, local_end_index - max_attention_size):local_end_index] + v_cache = kv_cache["v"][:, max(0, local_end_index - max_attention_size):local_end_index] x = attention(roped_query, k_cache, v_cache) kv_cache["global_end_index"].fill_(current_end) kv_cache["local_end_index"].fill_(local_end_index) @@ -364,6 +359,7 @@ class CausalWanSelfAttention(nn.Module): x = sequence_model_parallel_all_to_all_4D(x, scatter_dim=1, gather_dim=2) return self.o(x.flatten(2)) + class WanCrossAttention(CausalWanSelfAttention): """LingBot World 2 cross-attention with reusable text K/V cache.""" @@ -379,7 +375,8 @@ class WanCrossAttention(CausalWanSelfAttention): b, n, d = x.size(0), self.num_heads, self.head_dim q = self.norm_q(self.q(x)).view(b, -1, n, d) if crossattn_cache is not None: - is_first = crossattn_cache["is_init"].item() == 0 if cross_attn_first_call is None else cross_attn_first_call + is_first = crossattn_cache["is_init"].item( + ) == 0 if cross_attn_first_call is None else cross_attn_first_call if is_first: crossattn_cache["is_init"].fill_(1) k = self.norm_k(self.k(context)).view(b, -1, n, d) @@ -470,9 +467,7 @@ class CausalWanAttentionBlock(nn.Module): if dit_cond_dict is not None and "c2ws_plucker_emb" in dit_cond_dict: c2ws_plucker_emb = dit_cond_dict["c2ws_plucker_emb"] - c2ws_hidden_states = self.cam_injector_layer2( - F.silu(self.cam_injector_layer1(c2ws_plucker_emb)) - ) + c2ws_hidden_states = self.cam_injector_layer2(F.silu(self.cam_injector_layer1(c2ws_plucker_emb))) c2ws_hidden_states = c2ws_hidden_states + c2ws_plucker_emb x = (1.0 + self.cam_scale_layer(c2ws_hidden_states)) * x + self.cam_shift_layer(c2ws_hidden_states) @@ -567,21 +562,18 @@ class LingBotWorld2CausalFastTransformer3DModel(BaseDiT): nn.Linear(self.dim, self.dim), ) self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(self.dim, self.dim * 6)) - self.blocks = nn.ModuleList( - [ - CausalWanAttentionBlock( - self.dim, - self.ffn_dim, - self.num_heads, - self.local_attn_size, - self.sink_size, - self.qk_norm, - self.cross_attn_norm, - self.eps, - ) - for _ in range(self.num_layers) - ] - ) + self.blocks = nn.ModuleList([ + CausalWanAttentionBlock( + self.dim, + self.ffn_dim, + self.num_heads, + self.local_attn_size, + self.sink_size, + self.qk_norm, + self.cross_attn_norm, + self.eps, + ) for _ in range(self.num_layers) + ]) self.head = CausalHead(self.dim, self.out_dim, self.patch_size, self.eps) self.freqs: torch.Tensor | None = None self.init_weights() @@ -667,14 +659,12 @@ class LingBotWorld2CausalFastTransformer3DModel(BaseDiT): bt = t.size(0) t = t.flatten() e = self.time_embedding( - sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, padded_seq_len)).float() - ) + sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, padded_seq_len)).float()) e0 = self.time_projection(e).unflatten(2, (6, self.dim)) context_lens = None context = self.text_embedding( - torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context]) - ) + torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context])) if dit_cond_dict is not None and "c2ws_plucker_emb" in dit_cond_dict: c2ws_plucker_emb = dit_cond_dict["c2ws_plucker_emb"] @@ -685,14 +675,12 @@ class LingBotWorld2CausalFastTransformer3DModel(BaseDiT): c1=self.patch_size[0], c2=self.patch_size[1], c3=self.patch_size[2], - ) - for i in c2ws_plucker_emb + ) for i in c2ws_plucker_emb ] c2ws_plucker_emb = torch.cat(c2ws_plucker_emb, dim=1) c2ws_plucker_emb = self.patch_embedding_wancamctrl(c2ws_plucker_emb) - c2ws_hidden_states = self.c2ws_hidden_states_layer2( - F.silu(self.c2ws_hidden_states_layer1(c2ws_plucker_emb)) - ) + c2ws_hidden_states = self.c2ws_hidden_states_layer2(F.silu( + self.c2ws_hidden_states_layer1(c2ws_plucker_emb))) c2ws_plucker_emb = c2ws_plucker_emb + c2ws_hidden_states cam_len = c2ws_plucker_emb.size(1) if cam_len < padded_seq_len: @@ -748,7 +736,7 @@ class LingBotWorld2CausalFastTransformer3DModel(BaseDiT): c = self.out_dim out = [] for u, v in zip(x, grid_sizes.tolist(), strict=True): - u = u[: math.prod(v)].view(*v, *self.patch_size, c) + u = u[:math.prod(v)].view(*v, *self.patch_size, c) u = torch.einsum("fhwpqrc->cfphqwr", u) u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=True)]) out.append(u) diff --git a/fastvideo/models/dits/longcat.py b/fastvideo/models/dits/longcat.py index 238750f1c..603161b75 100644 --- a/fastvideo/models/dits/longcat.py +++ b/fastvideo/models/dits/longcat.py @@ -21,21 +21,21 @@ from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum from fastvideo.third_party.longcat_video.block_sparse_attention.bsa_interface import flash_attn_bsa_3d - # ============================================================================ # Embeddings # ============================================================================ + class PatchEmbed3D(nn.Module): """ 3D patch embedding using Conv3d. """ def __init__( - self, - patch_size: tuple[int, int, int] = (1, 2, 2), - in_channels: int = 16, - embed_dim: int = 4096, + self, + patch_size: tuple[int, int, int] = (1, 2, 2), + in_channels: int = 16, + embed_dim: int = 4096, ): super().__init__() self.patch_size = patch_size @@ -106,9 +106,8 @@ class TimestepEmbedder(nn.Module): Create sinusoidal timestep embeddings. """ half = dim // 2 - freqs = torch.exp( - -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half - ) + freqs = torch.exp(-math.log(max_period) * + torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half) args = t[:, None].float() * freqs[None] embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) if dim % 2: @@ -125,25 +124,25 @@ class TimestepEmbedder(nn.Module): """ # Sinusoidal embedding in FP32 t_freq = self.timestep_embedding(t.flatten(), self.frequency_embedding_size) - + # Cast to model dtype before MLP (matching original LongCat) # Handle LoRA wrapper if present linear_layer = self.linear_1.base_layer if hasattr(self.linear_1, 'base_layer') else self.linear_1 target_dtype = linear_layer.weight.dtype if t_freq.dtype != target_dtype: t_freq = t_freq.to(target_dtype) - + # MLP projection t_emb, _ = self.linear_1(t_freq) t_emb = self.act(t_emb) t_emb, _ = self.linear_2(t_emb) - + # Reshape if needed if latent_shape is not None and len(t.shape) > 1: B = t.shape[0] T = latent_shape[0] t_emb = t_emb.reshape(B, T, -1) - + return t_emb @@ -198,7 +197,7 @@ class CaptionEmbedder(nn.Module): y, _ = self.linear_1(encoder_hidden_states) y = self.act(y) y, _ = self.linear_2(y) # [B, N_text, C] - + # Handle attention masking - just zero out padded tokens if requested if encoder_attention_mask is not None: # Remove extra dimensions @@ -217,11 +216,11 @@ class CaptionEmbedder(nn.Module): ) elif mask_len > seq_len: encoder_attention_mask = encoder_attention_mask[:, :seq_len] - + # Zero out padded tokens if requested if self.text_tokens_zero_pad: y = y * encoder_attention_mask.unsqueeze(-1) - + # Return standard format [B, N_text, C] - no compaction! return y @@ -230,6 +229,7 @@ class CaptionEmbedder(nn.Module): # Attention Modules (Placeholders for now) # ============================================================================ + class LongCatSelfAttention(nn.Module): """ Self-attention with 3D RoPE support and optional BSA. @@ -274,13 +274,12 @@ class LongCatSelfAttention(nn.Module): ) def forward( - self, - x: torch.Tensor, # [B, N, C] - latent_shape: tuple, # (T, H, W) - num_cond_latents: int = 0, # Number of conditioning latent frames (for I2V) - return_kv: bool = False, # Return K/V for caching - **kwargs - ) -> torch.Tensor | tuple: + self, + x: torch.Tensor, # [B, N, C] + latent_shape: tuple, # (T, H, W) + num_cond_latents: int = 0, # Number of conditioning latent frames (for I2V) + return_kv: bool = False, # Return K/V for caching + **kwargs) -> torch.Tensor | tuple: """ Forward pass with 3D RoPE and optional BSA. @@ -330,41 +329,38 @@ class LongCatSelfAttention(nn.Module): if num_cond_latents > 0: # Calculate number of conditioned tokens (cond_latents * spatial_tokens_per_frame) num_cond_tokens = num_cond_latents * (N // T) - + # Conditioned tokens: only attend to themselves (same seq length, use self.attn) q_cond = q[:, :num_cond_tokens].contiguous() k_cond = k[:, :num_cond_tokens].contiguous() v_cond = v[:, :num_cond_tokens].contiguous() out_cond, _ = self.attn(q_cond, k_cond, v_cond) - + # Noise tokens: attend to ALL tokens (different seq lengths!) # Need to use flash attention directly since q has different length than k/v q_noise = q[:, num_cond_tokens:].contiguous() # [B, N_noise, num_heads, head_dim] # k, v are full: [B, N, num_heads, head_dim] - + # Transpose for flash attention: [B, num_heads, seq, head_dim] q_noise_t = q_noise.transpose(1, 2) k_t = k.transpose(1, 2) v_t = v.transpose(1, 2) - + # Use scaled dot product attention (handles different q/kv lengths) out_noise_t = torch.nn.functional.scaled_dot_product_attention( - q_noise_t, k_t, v_t, - attn_mask=None, - dropout_p=0.0, - is_causal=False - ) # [B, num_heads, N_noise, head_dim] - + q_noise_t, k_t, v_t, attn_mask=None, dropout_p=0.0, + is_causal=False) # [B, num_heads, N_noise, head_dim] + # Transpose back: [B, N_noise, num_heads, head_dim] out_noise = out_noise_t.transpose(1, 2) - + # Merge conditioned and noise outputs out = torch.cat([out_cond, out_noise], dim=1) - + # Reshape and project out out = out.reshape(B, N, C) out, _ = self.to_out(out) - + if return_kv: return out, (k_cache, v_cache) return out @@ -375,7 +371,7 @@ class LongCatSelfAttention(nn.Module): q_bsa = q.transpose(1, 2).contiguous() # [B, num_heads, N, head_dim] k_bsa = k.transpose(1, 2).contiguous() v_bsa = v.transpose(1, 2).contiguous() - + # Handle SP split: BSA operates on per-rank spatial dimensions # Replicate LongCat's cp_split_hw logic exactly from fastvideo.distributed.parallel_state import get_sp_world_size @@ -387,7 +383,7 @@ class LongCatSelfAttention(nn.Module): if sp_size % i == 0: factors.append([i, sp_size // i]) cp_split_hw = min(factors, key=lambda x: abs(x[0] - x[1])) - + # Split H and W dimensions by their respective factors T_bsa, H_bsa, W_bsa = latent_shape assert H_bsa % cp_split_hw[0] == 0 and W_bsa % cp_split_hw[1] == 0, \ @@ -397,15 +393,15 @@ class LongCatSelfAttention(nn.Module): latent_shape_bsa = (T_bsa, H_bsa, W_bsa) else: latent_shape_bsa = latent_shape - + # Call BSA with per-rank latent shape - out = flash_attn_bsa_3d( - q_bsa, k_bsa, v_bsa, - latent_shape_q=latent_shape_bsa, - latent_shape_k=latent_shape_bsa, - **self.bsa_params - ) # [B, num_heads, N, head_dim] - + out = flash_attn_bsa_3d(q_bsa, + k_bsa, + v_bsa, + latent_shape_q=latent_shape_bsa, + latent_shape_k=latent_shape_bsa, + **self.bsa_params) # [B, num_heads, N, head_dim] + # Transpose back: [B, N, num_heads, head_dim] out = out.transpose(1, 2) else: @@ -421,11 +417,11 @@ class LongCatSelfAttention(nn.Module): return out def forward_with_kv_cache( - self, - x: torch.Tensor, # [B, N_noise, C] - only noise tokens - latent_shape: tuple, # (T_noise, H, W) - shape for noise only - num_cond_latents: int, # Number of conditioning latent frames - kv_cache: tuple, # (k_cond, v_cond) - [B, heads, N_cond, head_dim] + self, + x: torch.Tensor, # [B, N_noise, C] - only noise tokens + latent_shape: tuple, # (T_noise, H, W) - shape for noise only + num_cond_latents: int, # Number of conditioning latent frames + kv_cache: tuple, # (k_cond, v_cond) - [B, heads, N_cond, head_dim] ) -> torch.Tensor: """ Forward using cached K/V from conditioning frames. @@ -476,30 +472,30 @@ class LongCatSelfAttention(nn.Module): # 1. Pad Q with dummy tokens matching k_cache shape # 2. Apply RoPE to full sequence (T_cond + T_noise) # 3. Extract only the noise portion of Q - + # Create dummy Q padding to fill conditioning positions # k_cache shape: [B, heads, N_cond, head_dim] q_padding = torch.cat([torch.empty_like(k_cache), q_rope], dim=2).contiguous() - + # Concatenate cached K with noise K for RoPE k_full = torch.cat([k_cache, k_rope], dim=2) v_full = torch.cat([v_cache, v], dim=2) - + # Apply RoPE to full sequence (includes both cond and noise positions) # Grid size: (T_cond + T_noise, H, W) full_T = num_cond_latents + T q_padding, k_full = self.rope_3d(q_padding, k_full, grid_size=(full_T, H, W)) - + # Extract only the noise portion of Q (last N tokens) q_rope = q_padding[:, :, -N:].contiguous() # Run attention: Q_noise attends to full K/V (cond + noise) - out = torch.nn.functional.scaled_dot_product_attention( - q_rope, k_full, v_full, - attn_mask=None, - dropout_p=0.0, - is_causal=False - ) # [B, heads, N_noise, head_dim] + out = torch.nn.functional.scaled_dot_product_attention(q_rope, + k_full, + v_full, + attn_mask=None, + dropout_p=0.0, + is_causal=False) # [B, heads, N_noise, head_dim] # Transpose back: [B, N_noise, heads, head_dim] out = out.transpose(1, 2) @@ -549,13 +545,12 @@ class LongCatCrossAttention(nn.Module): ) def forward( - self, - x: torch.Tensor, # [B, N_img, C] - context: torch.Tensor, # [B, N_text, C] - latent_shape: tuple = None, # (T, H, W) - needed for I2V - num_cond_latents: int = 0, # Number of conditioning latent frames (for I2V) - **kwargs - ) -> torch.Tensor: + self, + x: torch.Tensor, # [B, N_img, C] + context: torch.Tensor, # [B, N_text, C] + latent_shape: tuple = None, # (T, H, W) - needed for I2V + num_cond_latents: int = 0, # Number of conditioning latent frames (for I2V) + **kwargs) -> torch.Tensor: """ Forward pass for cross-attention (standard implementation). @@ -575,15 +570,15 @@ class LongCatCrossAttention(nn.Module): if num_cond_latents > 0 and latent_shape is not None: T, H, W = latent_shape num_cond_tokens = num_cond_latents * (N_img // T) - + # Only process noise tokens x_noise = x[:, num_cond_tokens:] # [B, N_noise, C] - + # Project Q, K, V for noise tokens only q, _ = self.to_q(x_noise) k, _ = self.to_k(context) v, _ = self.to_v(context) - + N_text = context.shape[1] N_noise = x_noise.shape[1] @@ -600,14 +595,10 @@ class LongCatCrossAttention(nn.Module): out_noise = self.attn(q, k, v) # [B, N_noise, num_heads, head_dim] out_noise = out_noise.reshape(B, N_noise, C) out_noise, _ = self.to_out(out_noise) - + # Conditioned tokens get zero output - out_cond = torch.zeros( - (B, num_cond_tokens, C), - dtype=out_noise.dtype, - device=out_noise.device - ) - + out_cond = torch.zeros((B, num_cond_tokens, C), dtype=out_noise.dtype, device=out_noise.device) + # Merge out = torch.cat([out_cond, out_noise], dim=1) return out @@ -617,7 +608,7 @@ class LongCatCrossAttention(nn.Module): q, _ = self.to_q(x) k, _ = self.to_k(context) v, _ = self.to_v(context) - + N_text = context.shape[1] # Reshape to heads @@ -644,6 +635,7 @@ class LongCatCrossAttention(nn.Module): # Feed-Forward Network # ============================================================================ + class LongCatSwiGLUFFN(nn.Module): """ SwiGLU feed-forward network using FastVideo's ReplicatedLinear. @@ -681,6 +673,7 @@ class LongCatSwiGLUFFN(nn.Module): # Modulation Utilities # ============================================================================ + def modulate_fp32(norm: nn.Module, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: """ Apply modulation in FP32 for numerical stability. @@ -688,15 +681,15 @@ def modulate_fp32(norm: nn.Module, x: torch.Tensor, shift: torch.Tensor, scale: Converts inputs to FP32 for the modulation operation, then casts back. """ orig_dtype = x.dtype - + # Convert to FP32 for numerical stability shift_fp32 = shift.float() scale_fp32 = scale.float() - + # Normalize and modulate in FP32 x_norm = norm(x.to(torch.float32)) x_mod = x_norm * (scale_fp32 + 1) + shift_fp32 - + return x_mod.to(orig_dtype) @@ -704,6 +697,7 @@ def modulate_fp32(norm: nn.Module, x: torch.Tensor, shift: torch.Tensor, scale: # Transformer Block # ============================================================================ + class LongCatTransformerBlock(nn.Module): """ Single-stream transformer block with: @@ -771,17 +765,16 @@ class LongCatTransformerBlock(nn.Module): ) def forward( - self, - x: torch.Tensor, # [B, N, C] - context: torch.Tensor, # [B, N_text, C] - t: torch.Tensor, # [B, T, C_t] - latent_shape: tuple, # (T, H, W) - num_cond_latents: int = 0, # Number of conditioning latent frames (for I2V) - return_kv: bool = False, # Return K/V for caching - kv_cache: tuple | None = None, # Pre-computed K/V cache - skip_crs_attn: bool = False, # Skip cross-attention (for cache init) - **kwargs - ) -> torch.Tensor | tuple: + self, + x: torch.Tensor, # [B, N, C] + context: torch.Tensor, # [B, N_text, C] + t: torch.Tensor, # [B, T, C_t] + latent_shape: tuple, # (T, H, W) + num_cond_latents: int = 0, # Number of conditioning latent frames (for I2V) + return_kv: bool = False, # Return K/V for caching + kv_cache: tuple | None = None, # Pre-computed K/V cache + skip_crs_attn: bool = False, # Skip cross-attention (for cache init) + **kwargs) -> torch.Tensor | tuple: """ Forward pass with AdaLN modulation. @@ -824,7 +817,7 @@ class LongCatTransformerBlock(nn.Module): kv_cache_new = None # Don't return cache when using cache else: attn_result = self.self_attn( - x_norm, + x_norm, latent_shape=latent_shape, num_cond_latents=num_cond_latents, return_kv=return_kv, @@ -845,12 +838,10 @@ class LongCatTransformerBlock(nn.Module): x_norm_cross = self.norm_cross(x) # When using KV cache, no need for num_cond_latents in cross-attn cross_num_cond = 0 if kv_cache is not None else num_cond_latents - cross_out = self.cross_attn( - x_norm_cross, - context, - latent_shape=latent_shape, - num_cond_latents=cross_num_cond - ) + cross_out = self.cross_attn(x_norm_cross, + context, + latent_shape=latent_shape, + num_cond_latents=cross_num_cond) x = x + cross_out # === FFN === @@ -873,6 +864,7 @@ class LongCatTransformerBlock(nn.Module): # Final Layer # ============================================================================ + class FinalLayer(nn.Module): """ Final output projection with AdaLN modulation. @@ -940,6 +932,7 @@ class FinalLayer(nn.Module): # Main Model # ============================================================================ + class LongCatTransformer3DModel(BaseDiT): """ Native LongCat Video Transformer using FastVideo layers. @@ -1004,8 +997,7 @@ class LongCatTransformer3DModel(BaseDiT): mlp_ratio=self.mlp_ratio, adaln_tembed_dim=config.adaln_tembed_dim, config=config, - ) - for _ in range(self.depth) + ) for _ in range(self.depth) ]) # Output projection @@ -1020,28 +1012,27 @@ class LongCatTransformer3DModel(BaseDiT): """Enable BSA for all self-attention layers.""" for block in self.blocks: block.self_attn.enable_bsa = True - + def disable_bsa(self): """Disable BSA for all self-attention layers.""" for block in self.blocks: block.self_attn.enable_bsa = False def forward( - self, - hidden_states: torch.Tensor, # [B, C, T, H, W] - encoder_hidden_states: torch.Tensor | list[torch.Tensor], # [B, N_text, C_text] - timestep: torch.LongTensor, # [B] or [B, T] - encoder_attention_mask: torch.Tensor | None = None, # [B, N_text] - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, - guidance: float | None = None, # Unused, for API compatibility - num_cond_latents: int = 0, # For I2V: number of conditioning latent frames - # === KV Cache Parameters === - return_kv: bool = False, # If True, return (output, kv_cache_dict) - kv_cache_dict: dict | None = None, # Pre-computed {block_idx: (k, v)} - skip_crs_attn: bool = False, # Skip cross-attention (for cache init) - offload_kv_cache: bool = False, # Move cache to CPU after compute - **kwargs - ) -> torch.Tensor | tuple[torch.Tensor, dict]: + self, + hidden_states: torch.Tensor, # [B, C, T, H, W] + encoder_hidden_states: torch.Tensor | list[torch.Tensor], # [B, N_text, C_text] + timestep: torch.LongTensor, # [B] or [B, T] + encoder_attention_mask: torch.Tensor | None = None, # [B, N_text] + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + guidance: float | None = None, # Unused, for API compatibility + num_cond_latents: int = 0, # For I2V: number of conditioning latent frames + # === KV Cache Parameters === + return_kv: bool = False, # If True, return (output, kv_cache_dict) + kv_cache_dict: dict | None = None, # Pre-computed {block_idx: (k, v)} + skip_crs_attn: bool = False, # Skip cross-attention (for cache init) + offload_kv_cache: bool = False, # Move cache to CPU after compute + **kwargs) -> torch.Tensor | tuple[torch.Tensor, dict]: """ Forward pass with FastVideo parameter ordering. @@ -1074,33 +1065,33 @@ class LongCatTransformer3DModel(BaseDiT): # Expand timestep from [B] to [B, T] if needed if timestep.ndim == 1: timestep = timestep.unsqueeze(1).expand(-1, N_t) # [B, T] - + t = self.time_embedder(timestep.flatten(), latent_shape=(N_t, N_h, N_w)) if t.ndim == 2: t = t.reshape(B, N_t, -1) # [B, T, C_t] # 3. Caption embedding (standard format, no compaction) - context = self.caption_embedder( - encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask - ) # [B, N_text, C] + context = self.caption_embedder(encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask) # [B, N_text, C] # 4. Transformer blocks with optional KV cache kv_cache_dict_ret = {} if return_kv else None - + for i, block in enumerate(self.blocks): # Get cache for this block if available block_kv_cache = kv_cache_dict.get(i, None) if kv_cache_dict else None - + block_out = block( - x, context, t, + x, + context, + t, latent_shape=(N_t, N_h, N_w), num_cond_latents=num_cond_latents, return_kv=return_kv, kv_cache=block_kv_cache, skip_crs_attn=skip_crs_attn, ) - + if return_kv: x, kv_cache = block_out # Store cache @@ -1145,5 +1136,6 @@ class LongCatTransformer3DModel(BaseDiT): ) return x + # Entry point for model registry EntryClass = LongCatTransformer3DModel diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index 24fd759d0..2ef569648 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -206,6 +206,7 @@ def get_timestep_embedding( class TimestepEmbedding(torch.nn.Module): """Two-layer MLP to project timestep embeddings.""" + def __init__( self, in_channels: int, @@ -237,6 +238,7 @@ class TimestepEmbedding(torch.nn.Module): class Timesteps(torch.nn.Module): """Sinusoidal timestep embedding wrapper with scaling knobs.""" + def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): super().__init__() self.num_channels = num_channels @@ -256,6 +258,7 @@ class Timesteps(torch.nn.Module): class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module): """PixArt-Alpha timestep embedding used by LTX-2 AdaLN.""" + def __init__(self, embedding_dim: int, size_emb_dim: int): super().__init__() self.outdim = size_emb_dim @@ -270,6 +273,7 @@ class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module): class AdaLayerNormSingle(torch.nn.Module): """AdaLN-single modulation that emits scale/shift/gates for LTX-2.""" + def __init__(self, embedding_dim: int, embedding_coefficient: int = 6): super().__init__() self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings( @@ -296,12 +300,12 @@ ADALN_NUM_CROSS_ATTN_PARAMS = 3 def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int: - return ADALN_NUM_BASE_PARAMS + (ADALN_NUM_CROSS_ATTN_PARAMS - if cross_attention_adaln else 0) + return ADALN_NUM_BASE_PARAMS + (ADALN_NUM_CROSS_ATTN_PARAMS if cross_attention_adaln else 0) class PixArtAlphaTextProjection(torch.nn.Module): """Caption projection MLP used by LTX-2.""" + def __init__(self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh"): super().__init__() if out_features is None: @@ -324,6 +328,7 @@ class PixArtAlphaTextProjection(torch.nn.Module): class GELUApprox(nn.Module): """Linear + tanh-approximate GELU used by LTX-2 FFN.""" + def __init__( self, in_features: int, @@ -346,6 +351,7 @@ class GELUApprox(nn.Module): class FeedForward(nn.Module): """LTX-2 FFN: GELUApprox -> Identity -> Linear.""" + def __init__( self, dim: int, @@ -379,6 +385,7 @@ class FeedForward(nn.Module): class VideoLatentShape(tuple): """Helper for (B, C, T, H, W) latent shapes.""" + @property def batch(self) -> int: return self[0] @@ -443,14 +450,13 @@ class AudioLatentShape(tuple): hop_length: int, audio_latent_downsample_factor: int, ) -> "AudioLatentShape": - latents_per_second = float(sample_rate) / float( - hop_length) / float(audio_latent_downsample_factor) - return AudioLatentShape( - (batch, channels, round(duration * latents_per_second), mel_bins)) + latents_per_second = float(sample_rate) / float(hop_length) / float(audio_latent_downsample_factor) + return AudioLatentShape((batch, channels, round(duration * latents_per_second), mel_bins)) class VideoLatentPatchifier: """Patchify/unpatchify latent tokens for LTX-2.""" + def __init__(self, patch_size: int): self._patch_size = (1, patch_size, patch_size) @@ -578,8 +584,7 @@ class AudioLatentPatchifier: torch.float32, device, ) - start_timings = start_timings.unsqueeze(0).expand(output_shape.batch, - -1).unsqueeze(1) + start_timings = start_timings.unsqueeze(0).expand(output_shape.batch, -1).unsqueeze(1) end_timings = self._get_audio_latent_time_in_sec( self.shift + 1, @@ -587,8 +592,7 @@ class AudioLatentPatchifier: torch.float32, device, ) - end_timings = end_timings.unsqueeze(0).expand(output_shape.batch, - -1).unsqueeze(1) + end_timings = end_timings.unsqueeze(0).expand(output_shape.batch, -1).unsqueeze(1) return torch.stack([start_timings, end_timings], dim=-1) @@ -600,14 +604,11 @@ class AudioLatentPatchifier: device: Optional[torch.device], ) -> torch.Tensor: resolved_device = device or torch.device("cpu") - audio_latent_frame = torch.arange( - start_latent, end_latent, dtype=dtype, device=resolved_device) + audio_latent_frame = torch.arange(start_latent, end_latent, dtype=dtype, device=resolved_device) audio_mel_frame = audio_latent_frame * self.audio_latent_downsample_factor if self.is_causal: causal_offset = 1 - audio_mel_frame = ( - audio_mel_frame + causal_offset - - self.audio_latent_downsample_factor).clamp(min=0) + audio_mel_frame = (audio_mel_frame + causal_offset - self.audio_latent_downsample_factor).clamp(min=0) return audio_mel_frame * self.hop_length / self.sample_rate @@ -626,8 +627,7 @@ def _get_pixel_coords( ).view(*broadcast_shape) pixel_coords = latent_coords.to(torch.float32) * scale_tensor if causal_fix: - pixel_coords[:, 0, ...] = ( - pixel_coords[:, 0, ...] + 1 - scale_factors[0]).clamp(min=0) + pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors[0]).clamp(min=0) if fps: pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps return pixel_coords @@ -673,23 +673,19 @@ def _debug_transformer_args(prefix: str, args: "TransformerArgs | None") -> None else: finite = torch.isfinite(mask) finite_sum = mask[finite].sum().item() if finite.any() else 0.0 - mask_summary = ( - f"mask_min={mask.min().item():.6f} " - f"mask_max={mask.max().item():.6f} " - f"mask_finite_sum={finite_sum:.6f} " - f"mask_finite_count={finite.sum().item()}" - ) - _debug_block_log_line( - f"{prefix}:x_sum={args.x.float().sum().item():.6f} " - f"context_sum={args.context.float().sum().item():.6f} " - f"t_sum={args.timesteps.float().sum().item():.6f} " - f"emb_sum={args.embedded_timestep.float().sum().item():.6f} " - f"pe_cos_sum={pe_cos.float().sum().item():.6f} " - f"pe_sin_sum={pe_sin.float().sum().item():.6f} " - f"cross_pe_cos_sum={(cross_cos.float().sum().item() if cross_cos is not None else 0.0):.6f} " - f"cross_pe_sin_sum={(cross_sin.float().sum().item() if cross_sin is not None else 0.0):.6f} " - f"{mask_summary}" - ) + mask_summary = (f"mask_min={mask.min().item():.6f} " + f"mask_max={mask.max().item():.6f} " + f"mask_finite_sum={finite_sum:.6f} " + f"mask_finite_count={finite.sum().item()}") + _debug_block_log_line(f"{prefix}:x_sum={args.x.float().sum().item():.6f} " + f"context_sum={args.context.float().sum().item():.6f} " + f"t_sum={args.timesteps.float().sum().item():.6f} " + f"emb_sum={args.embedded_timestep.float().sum().item():.6f} " + f"pe_cos_sum={pe_cos.float().sum().item():.6f} " + f"pe_sin_sum={pe_sin.float().sum().item():.6f} " + f"cross_pe_cos_sum={(cross_cos.float().sum().item() if cross_cos is not None else 0.0):.6f} " + f"cross_pe_sin_sum={(cross_sin.float().sum().item() if cross_sin is not None else 0.0):.6f} " + f"{mask_summary}") class LTXRopeType(Enum): @@ -759,9 +755,8 @@ def apply_ltx_rotary_emb_4d( raise ValueError(f"Invalid rope type: {rope_type}") -def _apply_ltx_interleaved_rotary_emb( - input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor -) -> torch.Tensor: +def _apply_ltx_interleaved_rotary_emb(input_tensor: torch.Tensor, cos_freqs: torch.Tensor, + sin_freqs: torch.Tensor) -> torch.Tensor: t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2) t1, t2 = t_dup.unbind(dim=-1) t_dup = torch.stack((-t2, t1), dim=-1) @@ -769,9 +764,8 @@ def _apply_ltx_interleaved_rotary_emb( return input_tensor * cos_freqs + input_tensor_rot * sin_freqs -def _apply_ltx_split_rotary_emb( - input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor -) -> torch.Tensor: +def _apply_ltx_split_rotary_emb(input_tensor: torch.Tensor, cos_freqs: torch.Tensor, + sin_freqs: torch.Tensor) -> torch.Tensor: needs_reshape = False if input_tensor.ndim != 4 and cos_freqs.ndim == 4: b, h, t, _ = cos_freqs.shape @@ -795,9 +789,8 @@ def _apply_ltx_split_rotary_emb( return output -def _generate_ltx_freq_grid_float64( - positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int -) -> torch.Tensor: +def _generate_ltx_freq_grid_float64(positional_embedding_theta: float, positional_embedding_max_pos_count: int, + inner_dim: int) -> torch.Tensor: theta = positional_embedding_theta start = 1 end = theta @@ -814,70 +807,60 @@ def _generate_ltx_freq_grid_float64( return (pow_indices * math.pi / 2).to(dtype=torch.float32) -_cached_generate_ltx_freq_grid_float64 = functools.lru_cache(maxsize=5)( - _generate_ltx_freq_grid_float64) +_cached_generate_ltx_freq_grid_float64 = functools.lru_cache(maxsize=5)(_generate_ltx_freq_grid_float64) -def generate_ltx_freq_grid_float64( - positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int -) -> torch.Tensor: +def generate_ltx_freq_grid_float64(positional_embedding_theta: float, positional_embedding_max_pos_count: int, + inner_dim: int) -> torch.Tensor: """Generate an eager-cached, compile-safe float64 LTX-2 rotary grid.""" if torch.compiler.is_compiling(): - return _generate_ltx_freq_grid_float64( - positional_embedding_theta, positional_embedding_max_pos_count, inner_dim) - return _cached_generate_ltx_freq_grid_float64( - positional_embedding_theta, positional_embedding_max_pos_count, inner_dim) + return _generate_ltx_freq_grid_float64(positional_embedding_theta, positional_embedding_max_pos_count, + inner_dim) + return _cached_generate_ltx_freq_grid_float64(positional_embedding_theta, positional_embedding_max_pos_count, + inner_dim) -def _generate_ltx_freq_grid_pytorch( - positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int -) -> torch.Tensor: +def _generate_ltx_freq_grid_pytorch(positional_embedding_theta: float, positional_embedding_max_pos_count: int, + inner_dim: int) -> torch.Tensor: theta = positional_embedding_theta start = 1 end = theta n_elem = 2 * positional_embedding_max_pos_count - indices = theta ** ( - torch.linspace( - math.log(start, theta), - math.log(end, theta), - inner_dim // n_elem, - dtype=torch.float32, - ) - ) + indices = theta**(torch.linspace( + math.log(start, theta), + math.log(end, theta), + inner_dim // n_elem, + dtype=torch.float32, + )) indices = indices.to(dtype=torch.float32) return indices * math.pi / 2 -_cached_generate_ltx_freq_grid_pytorch = functools.lru_cache(maxsize=5)( - _generate_ltx_freq_grid_pytorch) +_cached_generate_ltx_freq_grid_pytorch = functools.lru_cache(maxsize=5)(_generate_ltx_freq_grid_pytorch) -def generate_ltx_freq_grid_pytorch( - positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int -) -> torch.Tensor: +def generate_ltx_freq_grid_pytorch(positional_embedding_theta: float, positional_embedding_max_pos_count: int, + inner_dim: int) -> torch.Tensor: """Generate an eager-cached, compile-safe float32 LTX-2 rotary grid.""" if torch.compiler.is_compiling(): - return _generate_ltx_freq_grid_pytorch( - positional_embedding_theta, positional_embedding_max_pos_count, inner_dim) - return _cached_generate_ltx_freq_grid_pytorch( - positional_embedding_theta, positional_embedding_max_pos_count, inner_dim) + return _generate_ltx_freq_grid_pytorch(positional_embedding_theta, positional_embedding_max_pos_count, + inner_dim) + return _cached_generate_ltx_freq_grid_pytorch(positional_embedding_theta, positional_embedding_max_pos_count, + inner_dim) def _ltx_get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor: n_pos_dims = indices_grid.shape[1] if n_pos_dims != len(max_pos): - raise ValueError( - f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})" - ) + raise ValueError(f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})") return torch.stack( [indices_grid[:, i] / max_pos[i] for i in range(n_pos_dims)], dim=-1, ) -def _ltx_generate_freqs( - indices: torch.Tensor, indices_grid: torch.Tensor, max_pos: list[int], use_middle_indices_grid: bool -) -> torch.Tensor: +def _ltx_generate_freqs(indices: torch.Tensor, indices_grid: torch.Tensor, max_pos: list[int], + use_middle_indices_grid: bool) -> torch.Tensor: if use_middle_indices_grid: indices_grid_start, indices_grid_end = indices_grid[..., 0], indices_grid[..., 1] indices_grid = (indices_grid_start + indices_grid_end) / 2.0 @@ -890,9 +873,8 @@ def _ltx_generate_freqs( return freqs -def _ltx_split_freqs_cis( - freqs: torch.Tensor, pad_size: int, num_attention_heads: int -) -> tuple[torch.Tensor, torch.Tensor]: +def _ltx_split_freqs_cis(freqs: torch.Tensor, pad_size: int, + num_attention_heads: int) -> tuple[torch.Tensor, torch.Tensor]: cos_freq = freqs.cos() sin_freq = freqs.sin() @@ -983,6 +965,7 @@ class Modality: class TransformerArgsPreprocessor: """Prepare LTX-2 transformer inputs (patchify, AdaLN, rope).""" + def __init__( # noqa: PLR0913 self, patchify_proj: torch.nn.Linear, @@ -1050,8 +1033,7 @@ class TransformerArgsPreprocessor: if attention_mask is None or torch.is_floating_point(attention_mask): return attention_mask return (attention_mask - 1).to(x_dtype).reshape( - (attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) - ) * torch.finfo(x_dtype).max + (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])) * torch.finfo(x_dtype).max def _prepare_positional_embeddings( self, @@ -1086,8 +1068,8 @@ class TransformerArgsPreprocessor: if not latent.is_contiguous(): latent = latent.contiguous() x = self.patchify_proj(latent) - timestep, embedded_timestep = self._prepare_timestep( - modality.timesteps, x.shape[0], modality.latent.dtype, self.adaln) + timestep, embedded_timestep = self._prepare_timestep(modality.timesteps, x.shape[0], modality.latent.dtype, + self.adaln) prompt_timestep = None if self.prompt_adaln is not None and modality.sigma is not None: prompt_timestep, _ = self._prepare_timestep( @@ -1291,20 +1273,13 @@ class LTXDistributedAttention(DistributedAttention): use_vsa = self.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN if use_vsa: - assert ( - replicated_q is None and replicated_k is None - and replicated_v is None - ), "Replicated QKV is not supported for VSA now" + assert (replicated_q is None and replicated_k is None + and replicated_v is None), "Replicated QKV is not supported for VSA now" if gate_compress is None: - raise ValueError( - "gate_compress must be provided when using VIDEO_SPARSE_ATTN" - ) + raise ValueError("gate_compress must be provided when using VIDEO_SPARSE_ATTN") - qkvg = torch.cat([q, k, v, gate_compress], - dim=0) # [4*batch, seq_len, heads, head_dim] - qkvg = sequence_model_parallel_all_to_all_4D(qkvg, - scatter_dim=2, - gather_dim=1) + qkvg = torch.cat([q, k, v, gate_compress], dim=0) # [4*batch, seq_len, heads, head_dim] + qkvg = sequence_model_parallel_all_to_all_4D(qkvg, scatter_dim=2, gather_dim=1) pad_seq_len = qkvg.shape[1] - original_seq_len qkvg = qkvg[:, :original_seq_len, :, :] @@ -1316,8 +1291,7 @@ class LTXDistributedAttention(DistributedAttention): cos_local = cos[:, head_start:head_end, :original_seq_len, :] sin_local = sin[:, head_start:head_end, :original_seq_len, :] qk_part = qkvg[:batch_size * 2].transpose(1, 2) - qk_part = apply_ltx_rotary_emb_4d( - qk_part, (cos_local, sin_local), self.rope_type) + qk_part = apply_ltx_rotary_emb_4d(qk_part, (cos_local, sin_local), self.rope_type) qkvg[:batch_size * 2] = qk_part.transpose(1, 2) qkvg = self.attn_impl.preprocess_qkv(qkvg, ctx_attn_metadata) @@ -1327,9 +1301,7 @@ class LTXDistributedAttention(DistributedAttention): output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len)) - output = sequence_model_parallel_all_to_all_4D(output, - scatter_dim=1, - gather_dim=2) + output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2) return output, None # Stack QKV @@ -1366,13 +1338,10 @@ class LTXDistributedAttention(DistributedAttention): # Concatenate with replicated QKV if provided if replicated_q is not None: assert replicated_k is not None and replicated_v is not None - replicated_qkv = torch.cat( - [replicated_q, replicated_k, replicated_v], - dim=0) # [3, seq_len, num_heads, head_dim] + replicated_qkv = torch.cat([replicated_q, replicated_k, replicated_v], + dim=0) # [3, seq_len, num_heads, head_dim] heads_per_rank = num_heads // world_size - replicated_qkv = replicated_qkv[:, :, local_rank * - heads_per_rank:(local_rank + 1) * - heads_per_rank] + replicated_qkv = replicated_qkv[:, :, local_rank * heads_per_rank:(local_rank + 1) * heads_per_rank] qkv = torch.cat([qkv, replicated_qkv], dim=1) q, k, v = qkv.chunk(3, dim=0) @@ -1385,8 +1354,7 @@ class LTXDistributedAttention(DistributedAttention): split_idx = original_seq_len replicated_output = output[:, split_idx:] output = output[:, :split_idx] - replicated_output = sequence_model_parallel_all_gather( - replicated_output.contiguous(), dim=2) + replicated_output = sequence_model_parallel_all_gather(replicated_output.contiguous(), dim=2) # Apply backend-specific postprocess_output output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) @@ -1464,9 +1432,7 @@ class LTXLocalAttention(LocalAttention): if self.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: if gate_compress is None: - raise ValueError( - "gate_compress must be provided when using VIDEO_SPARSE_ATTN" - ) + raise ValueError("gate_compress must be provided when using VIDEO_SPARSE_ATTN") output = self.attn_impl.forward( # type: ignore[call-arg] q, k, v, gate_compress, ctx_attn_metadata) else: @@ -1476,6 +1442,7 @@ class LTXLocalAttention(LocalAttention): class LTXSelfAttention(nn.Module): """LTX-2 attention block with RMSNorm + FastVideo LocalAttention.""" + def __init__( self, query_dim: int, @@ -1570,7 +1537,7 @@ class LTXSelfAttention(nn.Module): dropout_rate=0, softmax_scale=None, causal=False, - supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA,), + supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA, ), default_backend=AttentionBackendEnum.TORCH_SDPA, ) @@ -1588,8 +1555,7 @@ class LTXSelfAttention(nn.Module): self.to_q.quant_method.quantize_input(x) # type: ignore[union-attr] if _supports_prequantized_input(self.to_q) else None) kv_pre_quantized = None - if (_supports_prequantized_input(self.to_k) - and _supports_prequantized_input(self.to_v)): + if (_supports_prequantized_input(self.to_k) and _supports_prequantized_input(self.to_v)): if context is x and q_pre_quantized is not None: kv_pre_quantized = q_pre_quantized else: @@ -1599,10 +1565,8 @@ class LTXSelfAttention(nn.Module): q = _linear_project_with_optional_prequant(self.to_q, x, q_pre_quantized) k = _linear_project_with_optional_prequant(self.to_k, context, kv_pre_quantized) v = _linear_project_with_optional_prequant(self.to_v, context, kv_pre_quantized) - gate_logits = (self.to_gate_logits(x)[0] - if self.to_gate_logits is not None else None) - gate_compress = (self.to_gate_compress(context)[0] - if self.to_gate_compress is not None else None) + gate_logits = (self.to_gate_logits(x)[0] if self.to_gate_logits is not None else None) + gate_compress = (self.to_gate_compress(context)[0] if self.to_gate_compress is not None else None) q = self.q_norm(q) k = self.k_norm(k) @@ -1614,8 +1578,7 @@ class LTXSelfAttention(nn.Module): k = k.view(b, k_len, self.heads, self.dim_head) v = v.view(b, k_len, self.heads, self.dim_head) if gate_compress is not None: - gate_compress = gate_compress.view(b, k_len, self.heads, - self.dim_head) + gate_compress = gate_compress.view(b, k_len, self.heads, self.dim_head) if mask is not None: if mask.ndim == 2: @@ -1634,22 +1597,13 @@ class LTXSelfAttention(nn.Module): attn_mask=mask, ) with set_forward_context( - current_timestep=current_timestep, - attn_metadata=attn_metadata, - forward_batch=forward_batch, + current_timestep=current_timestep, + attn_metadata=attn_metadata, + forward_batch=forward_batch, ): - out = self.attn_masked(q, - k, - v, - ltx_freqs_cis=pe, - ltx_k_freqs_cis=k_pe) + out = self.attn_masked(q, k, v, ltx_freqs_cis=pe, ltx_k_freqs_cis=k_pe) else: - out = self.attn(q, - k, - v, - gate_compress=gate_compress, - ltx_freqs_cis=pe, - ltx_k_freqs_cis=k_pe) + out = self.attn(q, k, v, gate_compress=gate_compress, ltx_freqs_cis=pe, ltx_k_freqs_cis=k_pe) # LTX-2.3 gated attention: scale the attention output per-head by # 2*sigmoid(gate_logits). No-op for LTX-2.0 (gate_logits is None). if gate_logits is not None: @@ -1773,8 +1727,7 @@ class LTXDistributedSelfAttention(nn.Module): self.to_q.quant_method.quantize_input(x) # type: ignore[union-attr] if _supports_prequantized_input(self.to_q) else None) kv_pre_quantized = None - if (_supports_prequantized_input(self.to_k) - and _supports_prequantized_input(self.to_v)): + if (_supports_prequantized_input(self.to_k) and _supports_prequantized_input(self.to_v)): if context is x and q_pre_quantized is not None: kv_pre_quantized = q_pre_quantized else: @@ -1784,10 +1737,8 @@ class LTXDistributedSelfAttention(nn.Module): q = _linear_project_with_optional_prequant(self.to_q, x, q_pre_quantized) k = _linear_project_with_optional_prequant(self.to_k, context, kv_pre_quantized) v = _linear_project_with_optional_prequant(self.to_v, context, kv_pre_quantized) - gate_logits = (self.to_gate_logits(x)[0] - if self.to_gate_logits is not None else None) - gate_compress = (self.to_gate_compress(context)[0] - if self.to_gate_compress is not None else None) + gate_logits = (self.to_gate_logits(x)[0] if self.to_gate_logits is not None else None) + gate_compress = (self.to_gate_compress(context)[0] if self.to_gate_compress is not None else None) q = self.q_norm(q) k = self.k_norm(k) @@ -1800,8 +1751,7 @@ class LTXDistributedSelfAttention(nn.Module): k = k.view(b, k_len, self.heads, self.dim_head) v = v.view(b, k_len, self.heads, self.dim_head) if gate_compress is not None: - gate_compress = gate_compress.view(b, k_len, self.heads, - self.dim_head) + gate_compress = gate_compress.view(b, k_len, self.heads, self.dim_head) # Pass full RoPE to distributed attention - it will apply after all-to-all # and slice to this rank's heads @@ -1860,8 +1810,7 @@ class BasicAVTransformerBlock(torch.nn.Module): AttentionBackendEnum.ATTN_QAT_INFER, AttentionBackendEnum.ATTN_QAT_TRAIN, ) - video_self_attn_backends = video_attn_backends + ( - AttentionBackendEnum.VIDEO_SPARSE_ATTN, ) + video_self_attn_backends = video_attn_backends + (AttentionBackendEnum.VIDEO_SPARSE_ATTN, ) audio_attn_backends = ( AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, @@ -1903,8 +1852,7 @@ class BasicAVTransformerBlock(torch.nn.Module): # 6 rows for LTX-2.0; 9 rows when cross_attention_adaln adds the # text cross-attention shift/scale/gate. video_sst_size = adaln_embedding_coefficient(cross_attention_adaln) - self.scale_shift_table = torch.nn.Parameter( - torch.empty(video_sst_size, video.dim)) + self.scale_shift_table = torch.nn.Parameter(torch.empty(video_sst_size, video.dim)) if audio is not None: # Audio self-attention - use distributed when SP > 1 @@ -1940,8 +1888,7 @@ class BasicAVTransformerBlock(torch.nn.Module): prefix=f"{prefix}.blocks.{idx}.audio", ) audio_sst_size = adaln_embedding_coefficient(cross_attention_adaln) - self.audio_scale_shift_table = torch.nn.Parameter( - torch.empty(audio_sst_size, audio.dim)) + self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(audio_sst_size, audio.dim)) if audio is not None and video is not None: # Audio-to-video cross-attention @@ -1978,11 +1925,9 @@ class BasicAVTransformerBlock(torch.nn.Module): # LTX-2.3 cross-attention AdaLN prompt scale/shift tables (only created # when cross_attention_adaln is enabled, so absent for LTX-2.0). if self.cross_attention_adaln and video is not None: - self.prompt_scale_shift_table = torch.nn.Parameter( - torch.empty(2, video.dim)) + self.prompt_scale_shift_table = torch.nn.Parameter(torch.empty(2, video.dim)) if self.cross_attention_adaln and audio is not None: - self.audio_prompt_scale_shift_table = torch.nn.Parameter( - torch.empty(2, audio.dim)) + self.audio_prompt_scale_shift_table = torch.nn.Parameter(torch.empty(2, audio.dim)) self.norm_eps = norm_eps @@ -1994,26 +1939,26 @@ class BasicAVTransformerBlock(torch.nn.Module): """ if not hasattr(self, 'unshard'): return # Not wrapped by FSDP2 - + def make_unshard_hook(): + def hook(grad): self.unshard() return grad + return hook - + if vx is not None and vx.requires_grad: vx.register_hook(make_unshard_hook()) if ax is not None and ax.requires_grad: ax.register_hook(make_unshard_hook()) - def get_ada_values( - self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice - ) -> tuple[torch.Tensor, ...]: + def get_ada_values(self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, + indices: slice) -> tuple[torch.Tensor, ...]: num_ada_params = scale_shift_table.shape[0] ada_values = ( - scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype) - + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :] - ).unbind(dim=2) + scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype) + + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :]).unbind(dim=2) return ada_values def get_av_ca_ada_values( @@ -2024,12 +1969,10 @@ class BasicAVTransformerBlock(torch.nn.Module): gate_timestep: torch.Tensor, num_scale_shift_values: int = 4, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - scale_shift_ada_values = self.get_ada_values( - scale_shift_table[:num_scale_shift_values, :], batch_size, scale_shift_timestep, slice(None, None) - ) - gate_ada_values = self.get_ada_values( - scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, slice(None, None) - ) + scale_shift_ada_values = self.get_ada_values(scale_shift_table[:num_scale_shift_values, :], batch_size, + scale_shift_timestep, slice(None, None)) + gate_ada_values = self.get_ada_values(scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, + slice(None, None)) scale_shift_chunks = [t.squeeze(2) for t in scale_shift_ada_values] gate_ada_values = [t.squeeze(2) for t in gate_ada_values] @@ -2054,8 +1997,7 @@ class BasicAVTransformerBlock(torch.nn.Module): """ if self.cross_attention_adaln: # The 3 extra AdaLN rows (6..9) carry text cross-attn shift/scale/gate. - shift_q, scale_q, gate = self.get_ada_values( - scale_shift_table, x.shape[0], timestep, slice(6, 9)) + shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x.shape[0], timestep, slice(6, 9)) return apply_cross_attention_adaln( x=x, context=context, @@ -2135,9 +2077,8 @@ class BasicAVTransformerBlock(torch.nn.Module): perturb = skip_flag.reshape(1).expand(bsz) else: if skip_flag.shape[0] != bsz: - raise ValueError( - "Per-sample STG mask batch size mismatch: " - f"got {skip_flag.shape[0]}, expected {bsz}") + raise ValueError("Per-sample STG mask batch size mismatch: " + f"got {skip_flag.shape[0]}, expected {bsz}") perturb = skip_flag.reshape(bsz) keep = 1.0 - perturb.to(device=values.device, dtype=values.dtype) elif bool(skip_flag): @@ -2145,11 +2086,9 @@ class BasicAVTransformerBlock(torch.nn.Module): return keep.view(bsz, *([1] * (values.ndim - 1))) if run_vx: - vshift_msa, vscale_msa, vgate_msa = self.get_ada_values( - self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3) - ) - norm_vx = _rms_norm_dispatch(vx, eps=self.norm_eps) * ( - 1 + vscale_msa) + vshift_msa + vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(self.scale_shift_table, vx.shape[0], + video.timesteps, slice(0, 3)) + norm_vx = _rms_norm_dispatch(vx, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa video_self_attn_mask = _build_attn_keep_mask(vx, skip_video_self_attn) if self.use_distributed_attention: vx_attn1_out = self.attn1( @@ -2166,19 +2105,16 @@ class BasicAVTransformerBlock(torch.nn.Module): context=video.context, attn=self.attn2, scale_shift_table=self.scale_shift_table, - prompt_scale_shift_table=getattr( - self, "prompt_scale_shift_table", None), + prompt_scale_shift_table=getattr(self, "prompt_scale_shift_table", None), timestep=video.timesteps, prompt_timestep=video.prompt_timestep, context_mask=video.context_mask, ) if run_ax: - ashift_msa, ascale_msa, agate_msa = self.get_ada_values( - self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3) - ) - norm_ax = _rms_norm_dispatch(ax, eps=self.norm_eps) * ( - 1 + ascale_msa) + ashift_msa + ashift_msa, ascale_msa, agate_msa = self.get_ada_values(self.audio_scale_shift_table, ax.shape[0], + audio.timesteps, slice(0, 3)) + norm_ax = _rms_norm_dispatch(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa audio_self_attn_mask = _build_attn_keep_mask(ax, skip_audio_self_attn) if self.use_distributed_attention: ax_attn1_out = self.audio_attn1( @@ -2195,8 +2131,7 @@ class BasicAVTransformerBlock(torch.nn.Module): context=audio.context, attn=self.audio_attn2, scale_shift_table=self.audio_scale_shift_table, - prompt_scale_shift_table=getattr( - self, "audio_prompt_scale_shift_table", None), + prompt_scale_shift_table=getattr(self, "audio_prompt_scale_shift_table", None), timestep=audio.timesteps, prompt_timestep=audio.prompt_timestep, context_mask=audio.context_mask, @@ -2254,28 +2189,22 @@ class BasicAVTransformerBlock(torch.nn.Module): if audio_k_pe is not None: original_audio_len = audio_k_pe[0].shape[2] ax_context = ax_context[:, :original_audio_len, :] - vx = vx + ( - self.audio_to_video_attn( - vx_scaled, - context=ax_context, - pe=video_q_pe, - k_pe=audio_k_pe, - ) - * gate_out_a2v - ) + vx = vx + (self.audio_to_video_attn( + vx_scaled, + context=ax_context, + pe=video_q_pe, + k_pe=audio_k_pe, + ) * gate_out_a2v) else: ax_context = ax_scaled video_q_pe = video.cross_positional_embeddings audio_k_pe = audio.cross_positional_embeddings - vx = vx + ( - self.audio_to_video_attn( - vx_scaled, - context=ax_context, - pe=video_q_pe, - k_pe=audio_k_pe, - ) - * gate_out_a2v - ) + vx = vx + (self.audio_to_video_attn( + vx_scaled, + context=ax_context, + pe=video_q_pe, + k_pe=audio_k_pe, + ) * gate_out_a2v) if run_v2a: ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a @@ -2311,33 +2240,26 @@ class BasicAVTransformerBlock(torch.nn.Module): vx_context = vx_scaled audio_q_pe = audio.cross_positional_embeddings video_k_pe = video.cross_positional_embeddings - ax = ax + ( - self.video_to_audio_attn( - ax_scaled, - context=vx_context, - pe=audio_q_pe, - k_pe=video_k_pe, - ) - * gate_out_v2a - ) + ax = ax + (self.video_to_audio_attn( + ax_scaled, + context=vx_context, + pe=audio_q_pe, + k_pe=video_k_pe, + ) * gate_out_v2a) if run_vx: # slice(3, 6) selects the FFN shift/scale/gate. Identical to # slice(3, None) for the 6-row LTX-2.0 table, but correct for the # 9-row cross_attention_adaln table (rows 6..9 are the cross-attn). - vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values( - self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, 6) - ) - vx_scaled = _rms_norm_dispatch(vx, eps=self.norm_eps) * ( - 1 + vscale_mlp) + vshift_mlp + vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(self.scale_shift_table, vx.shape[0], + video.timesteps, slice(3, 6)) + vx_scaled = _rms_norm_dispatch(vx, eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp vx = vx + self.ff(vx_scaled) * vgate_mlp if run_ax: - ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values( - self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, 6) - ) - ax_scaled = _rms_norm_dispatch(ax, eps=self.norm_eps) * ( - 1 + ascale_mlp) + ashift_mlp + ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(self.audio_scale_shift_table, ax.shape[0], + audio.timesteps, slice(3, 6)) + ax_scaled = _rms_norm_dispatch(ax, eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp ax = ax + self.audio_ff(ax_scaled) * agate_mlp # Debug-only: reading ``self.idx`` (and .item() syncs) here means @@ -2346,10 +2268,8 @@ class BasicAVTransformerBlock(torch.nn.Module): if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": video_sum = vx.float().sum().item() if vx is not None else 0.0 audio_sum = ax.float().sum().item() if ax is not None else 0.0 - _debug_block_log_line( - f"fastvideo:block={self.idx}:video_sum={video_sum:.6f} " - f"audio_sum={audio_sum:.6f}" - ) + _debug_block_log_line(f"fastvideo:block={self.idx}:video_sum={video_sum:.6f} " + f"audio_sum={audio_sum:.6f}") # Register FSDP2 backward hooks on output tensors (module-level hooks don't # fire for dataclass outputs, so we must hook the tensors directly) @@ -2380,18 +2300,13 @@ def apply_cross_attention_adaln( affects the LTX-2.0 path. """ if prompt_scale_shift_table is None or prompt_timestep is None: - raise ValueError( - "cross_attention_adaln requires prompt scale/shift tables and prompt_timestep." - ) + raise ValueError("cross_attention_adaln requires prompt scale/shift tables and prompt_timestep.") batch_size = x.shape[0] - shift_kv, scale_kv = ( - prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) - + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1) - ).unbind(dim=2) + shift_kv, scale_kv = (prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)).unbind(dim=2) attn_input = _rms_norm_dispatch(x, eps=norm_eps) * (1 + q_scale) + q_shift encoder_hidden_states = context * (1 + scale_kv) + shift_kv - return attn(attn_input, context=encoder_hidden_states, - mask=context_mask) * q_gate + return attn(attn_input, context=encoder_hidden_states, mask=context_mask) * q_gate class LTXModelType(Enum): @@ -2516,8 +2431,7 @@ class LTXModel(torch.nn.Module): self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True) self.adaln_single = AdaLayerNormSingle( self.inner_dim, - embedding_coefficient=adaln_embedding_coefficient( - self.cross_attention_adaln), + embedding_coefficient=adaln_embedding_coefficient(self.cross_attention_adaln), ) # When caption_proj_before_connector is set the caption projection # lives in the text encoder, so it is omitted here (LTX-2.3 layout). @@ -2546,8 +2460,7 @@ class LTXModel(torch.nn.Module): self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True) self.audio_adaln_single = AdaLayerNormSingle( self.audio_inner_dim, - embedding_coefficient=adaln_embedding_coefficient( - self.cross_attention_adaln), + embedding_coefficient=adaln_embedding_coefficient(self.cross_attention_adaln), ) if create_caption_projection: self.audio_caption_projection = PixArtAlphaTextProjection( @@ -2664,48 +2577,37 @@ class LTXModel(torch.nn.Module): quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: - video_config = ( - TransformerConfig( - dim=self.inner_dim, - heads=self.num_attention_heads, - d_head=attention_head_dim, - context_dim=cross_attention_dim, - apply_gated_attention=self.apply_gated_attention, - cross_attention_adaln=self.cross_attention_adaln, - ) - if self.model_type.is_video_enabled() - else None - ) - audio_config = ( - TransformerConfig( - dim=self.audio_inner_dim, - heads=self.audio_num_attention_heads, - d_head=audio_attention_head_dim, - context_dim=audio_cross_attention_dim, - apply_gated_attention=self.apply_gated_attention, - cross_attention_adaln=self.cross_attention_adaln, - ) - if self.model_type.is_audio_enabled() - else None - ) + video_config = (TransformerConfig( + dim=self.inner_dim, + heads=self.num_attention_heads, + d_head=attention_head_dim, + context_dim=cross_attention_dim, + apply_gated_attention=self.apply_gated_attention, + cross_attention_adaln=self.cross_attention_adaln, + ) if self.model_type.is_video_enabled() else None) + audio_config = (TransformerConfig( + dim=self.audio_inner_dim, + heads=self.audio_num_attention_heads, + d_head=audio_attention_head_dim, + context_dim=audio_cross_attention_dim, + apply_gated_attention=self.apply_gated_attention, + cross_attention_adaln=self.cross_attention_adaln, + ) if self.model_type.is_audio_enabled() else None) self.use_distributed_attention = use_distributed_attention - self.transformer_blocks = torch.nn.ModuleList( - [ - BasicAVTransformerBlock( - idx=idx, - video=video_config, - audio=audio_config, - rope_type=self.rope_type, - norm_eps=norm_eps, - use_distributed_attention=use_distributed_attention, - cross_attention_adaln=self.cross_attention_adaln, - stg_block_idx=self.stg_block_idx, - quant_config=quant_config, - prefix=prefix, - ) - for idx in range(num_layers) - ] - ) + self.transformer_blocks = torch.nn.ModuleList([ + BasicAVTransformerBlock( + idx=idx, + video=video_config, + audio=audio_config, + rope_type=self.rope_type, + norm_eps=norm_eps, + use_distributed_attention=use_distributed_attention, + cross_attention_adaln=self.cross_attention_adaln, + stg_block_idx=self.stg_block_idx, + quant_config=quant_config, + prefix=prefix, + ) for idx in range(num_layers) + ]) def _process_transformer_blocks( self, @@ -2745,9 +2647,8 @@ class LTXModel(torch.nn.Module): x: torch.Tensor, embedded_timestep: torch.Tensor, ) -> torch.Tensor: - scale_shift_values = ( - scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + embedded_timestep[:, :, None] - ) + scale_shift_values = (scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + + embedded_timestep[:, :, None]) shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] x = norm_out(x) x = x * (1 + scale) + shift @@ -2780,13 +2681,11 @@ class LTXModel(torch.nn.Module): self-attention is skipped (STG perturbed pass). """ if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": - _debug_block_log_line( - "fastvideo:patchify_proj" - f":video_w_sum={self.patchify_proj.weight.float().sum().item():.6f} " - f"video_b_sum={self.patchify_proj.bias.float().sum().item():.6f} " - f"audio_w_sum={self.audio_patchify_proj.weight.float().sum().item():.6f} " - f"audio_b_sum={self.audio_patchify_proj.bias.float().sum().item():.6f}" - ) + _debug_block_log_line("fastvideo:patchify_proj" + f":video_w_sum={self.patchify_proj.weight.float().sum().item():.6f} " + f"video_b_sum={self.patchify_proj.bias.float().sum().item():.6f} " + f"audio_w_sum={self.audio_patchify_proj.weight.float().sum().item():.6f} " + f"audio_b_sum={self.audio_patchify_proj.bias.float().sum().item():.6f}") if not self.model_type.is_video_enabled() and video is not None: raise ValueError("Video is not enabled for this model") if not self.model_type.is_audio_enabled() and audio is not None: @@ -2806,24 +2705,15 @@ class LTXModel(torch.nn.Module): skip_audio_self_attn_blocks=skip_audio_self_attn_blocks, ) - vx = ( - self._process_output( - self.scale_shift_table, self.norm_out, self.proj_out, video_out.x, video_out.embedded_timestep - ) - if video_out is not None - else None - ) - ax = ( - self._process_output( - self.audio_scale_shift_table, - self.audio_norm_out, - self.audio_proj_out, - audio_out.x, - audio_out.embedded_timestep, - ) - if audio_out is not None - else None - ) + vx = (self._process_output(self.scale_shift_table, self.norm_out, self.proj_out, video_out.x, + video_out.embedded_timestep) if video_out is not None else None) + ax = (self._process_output( + self.audio_scale_shift_table, + self.audio_norm_out, + self.audio_proj_out, + audio_out.x, + audio_out.embedded_timestep, + ) if audio_out is not None else None) return vx, ax @@ -2845,27 +2735,21 @@ class LTX2Transformer3DModel(BaseDiT): # Get SP world size for distributed attention sp_world_size = get_sp_world_size() - use_vsa_backend = os.getenv("FASTVIDEO_ATTENTION_BACKEND", - "") == "VIDEO_SPARSE_ATTN" + use_vsa_backend = os.getenv("FASTVIDEO_ATTENTION_BACKEND", "") == "VIDEO_SPARSE_ATTN" use_distributed_attention = sp_world_size > 1 or use_vsa_backend # Validate that attention heads are divisible by SP world size if sp_world_size > 1: assert arch.num_attention_heads % sp_world_size == 0, ( f"The number of video attention heads ({arch.num_attention_heads}) " - f"must be divisible by the sequence parallel size ({sp_world_size})" - ) + f"must be divisible by the sequence parallel size ({sp_world_size})") assert arch.audio_num_attention_heads % sp_world_size == 0, ( f"The number of audio attention heads ({arch.audio_num_attention_heads}) " - f"must be divisible by the sequence parallel size ({sp_world_size})" - ) - logger.info( - f"LTX2 sequence parallelism enabled with SP world size {sp_world_size}" - ) + f"must be divisible by the sequence parallel size ({sp_world_size})") + logger.info(f"LTX2 sequence parallelism enabled with SP world size {sp_world_size}") elif use_vsa_backend: - logger.info( - "LTX2 VSA enabled with SP world size 1; using distributed " - "attention path for VSA compatibility") + logger.info("LTX2 VSA enabled with SP world size 1; using distributed " + "attention path for VSA compatibility") model_type = LTXModelType.AudioVideo self.model = LTXModel( @@ -2931,42 +2815,44 @@ class LTX2Transformer3DModel(BaseDiT): return f"{tensor.float().sum().item():.6f}" def _hook_factory(block_idx: int, name: str): + def _hook(_module, _inputs, outputs): # noqa: ANN001 out = outputs[0] if isinstance(outputs, tuple) else outputs out_sum = _format_sum(out if torch.is_tensor(out) else None) with path.open("a", encoding="utf-8") as f: f.write(f"fastvideo:{block_idx}:{name}:out_sum={out_sum}\n") + return _hook for block in self.model.transformer_blocks: idx = block.idx for name in ( - "attn1", - "attn2", - "ff", - "audio_attn1", - "audio_attn2", - "audio_ff", - "audio_to_video_attn", - "video_to_audio_attn", + "attn1", + "attn2", + "ff", + "audio_attn1", + "audio_attn2", + "audio_ff", + "audio_to_video_attn", + "video_to_audio_attn", ): if hasattr(block, name): - getattr(block, name).register_forward_hook( - _hook_factory(idx, name)) + getattr(block, name).register_forward_hook(_hook_factory(idx, name)) def _output_hook(label: str): + def _hook(_module, _inputs, outputs): # noqa: ANN001 out = outputs[0] if isinstance(outputs, tuple) else outputs out_sum = _format_sum(out if torch.is_tensor(out) else None) with path.open("a", encoding="utf-8") as f: f.write(f"fastvideo:output:{label}:out_sum={out_sum}\n") + return _hook if hasattr(self.model, "proj_out"): self.model.proj_out.register_forward_hook(_output_hook("proj_out")) if hasattr(self.model, "audio_proj_out"): - self.model.audio_proj_out.register_forward_hook( - _output_hook("audio_proj_out")) + self.model.audio_proj_out.register_forward_hook(_output_hook("audio_proj_out")) def forward( self, @@ -3030,9 +2916,7 @@ class LTX2Transformer3DModel(BaseDiT): # This is necessary because sequence sharding may not align to frame boundaries. # The full RoPE will be passed to DistributedAttention which applies it after # the all-to-all (when each rank has the full sequence but subset of heads). - positions = self.patchifier.get_patch_grid_bounds( - video_shape, device=hidden_states.device - ) + positions = self.patchifier.get_patch_grid_bounds(video_shape, device=hidden_states.device) positions = _get_pixel_coords( positions, DEFAULT_LTX2_SCALE_FACTORS, @@ -3062,15 +2946,13 @@ class LTX2Transformer3DModel(BaseDiT): video_head = latents.flatten()[:8].float().tolist() video_flat = latents.float().flatten() video_checksum = (video_flat * torch.arange(video_flat.numel(), device=video_flat.device)).sum().item() - _debug_block_log_line( - "fastvideo:modality_video" - f":latent_sum={latents.float().sum().item():.6f} " - f"latent_shape={tuple(latents.shape)} " - f"positions_sum={positions.float().sum().item():.6f} " - f"positions_shape={tuple(positions.shape)} " - f"latent_head={video_head} " - f"latent_checksum={video_checksum:.6f}" - ) + _debug_block_log_line("fastvideo:modality_video" + f":latent_sum={latents.float().sum().item():.6f} " + f"latent_shape={tuple(latents.shape)} " + f"positions_sum={positions.float().sum().item():.6f} " + f"positions_shape={tuple(positions.shape)} " + f"latent_head={video_head} " + f"latent_checksum={video_checksum:.6f}") # Process audio modality if provided audio_modality = None @@ -3081,10 +2963,7 @@ class LTX2Transformer3DModel(BaseDiT): audio_shape = AudioLatentShape.from_torch_shape(audio_hidden_states.shape) audio_latents = self.audio_patchifier.patchify(audio_hidden_states) if audio_sigma is None: - audio_sigma = ( - audio_timestep[:, 0] - if audio_timestep.ndim > 1 else audio_timestep - ) + audio_sigma = (audio_timestep[:, 0] if audio_timestep.ndim > 1 else audio_timestep) # Shard audio latents and timestep across SP ranks audio_original_seq_len = audio_latents.shape[1] @@ -3096,9 +2975,8 @@ class LTX2Transformer3DModel(BaseDiT): # Compute audio RoPE positions for the FULL sequence (before sharding) # Same as video: full RoPE is applied after all-to-all in DistributedAttention - audio_positions = self.audio_patchifier.get_patch_grid_bounds( - audio_shape, device=audio_hidden_states.device - ) + audio_positions = self.audio_patchifier.get_patch_grid_bounds(audio_shape, + device=audio_hidden_states.device) audio_modality = Modality( enabled=True, @@ -3113,15 +2991,13 @@ class LTX2Transformer3DModel(BaseDiT): audio_head = audio_latents.flatten()[:8].float().tolist() audio_flat = audio_latents.float().flatten() audio_checksum = (audio_flat * torch.arange(audio_flat.numel(), device=audio_flat.device)).sum().item() - _debug_block_log_line( - "fastvideo:modality_audio" - f":latent_sum={audio_latents.float().sum().item():.6f} " - f"latent_shape={tuple(audio_latents.shape)} " - f"positions_sum={audio_positions.float().sum().item():.6f} " - f"positions_shape={tuple(audio_positions.shape)} " - f"latent_head={audio_head} " - f"latent_checksum={audio_checksum:.6f}" - ) + _debug_block_log_line("fastvideo:modality_audio" + f":latent_sum={audio_latents.float().sum().item():.6f} " + f"latent_shape={tuple(audio_latents.shape)} " + f"positions_sum={audio_positions.float().sum().item():.6f} " + f"positions_shape={tuple(audio_positions.shape)} " + f"latent_head={audio_head} " + f"latent_checksum={audio_checksum:.6f}") # Run transformer with original sequence lengths video_out, audio_out = self.model( @@ -3150,24 +3026,19 @@ class LTX2Transformer3DModel(BaseDiT): # Gather and unpad video output if sp_world_size > 1 and video_out is not None: - video_out = sequence_model_parallel_all_gather_with_unpad( - video_out, video_original_seq_len, dim=1 - ) + video_out = sequence_model_parallel_all_gather_with_unpad(video_out, video_original_seq_len, dim=1) # Gather and unpad audio output if sp_world_size > 1 and audio_out is not None: - audio_out = sequence_model_parallel_all_gather_with_unpad( - audio_out, audio_original_seq_len, dim=1 - ) + audio_out = sequence_model_parallel_all_gather_with_unpad(audio_out, audio_original_seq_len, dim=1) # Unpatchify - video_out = self.patchifier.unpatchify( - video_out, output_shape=video_shape) + video_out = self.patchifier.unpatchify(video_out, output_shape=video_shape) if audio_out is None or audio_shape is None: return video_out - audio_out = self.audio_patchifier.unpatchify( - audio_out, output_shape=audio_shape) + audio_out = self.audio_patchifier.unpatchify(audio_out, output_shape=audio_shape) return video_out, audio_out + # Entry point for model registry EntryClass = LTX2Transformer3DModel diff --git a/fastvideo/models/dits/magi_human.py b/fastvideo/models/dits/magi_human.py index 4682326bf..88d7d13bf 100644 --- a/fastvideo/models/dits/magi_human.py +++ b/fastvideo/models/dits/magi_human.py @@ -59,7 +59,6 @@ from fastvideo.layers.rotary_embedding import _apply_rotary_emb from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum - # --------------------------------------------------------------------------- # Enums # --------------------------------------------------------------------------- @@ -167,9 +166,7 @@ class MultiModalityRMSNorm(nn.Module): t = self._rms(x) if self.num_modality == 1: return (t * (self.weight + 1)).to(original_dtype) - assert modality_dispatcher is not None, ( - "MultiModalityRMSNorm with num_modality>1 requires a dispatcher" - ) + assert modality_dispatcher is not None, ("MultiModalityRMSNorm with num_modality>1 requires a dispatcher") weight_chunks = self.weight.chunk(self.num_modality, dim=0) parts = modality_dispatcher.dispatch(t) for i in range(self.num_modality): @@ -179,7 +176,7 @@ class MultiModalityRMSNorm(nn.Module): def _freq_bands(num_bands: int, temperature: float = 10000.0) -> torch.Tensor: exp = torch.arange(0, num_bands, 1, dtype=torch.int64).float() / num_bands - return 1.0 / (temperature ** exp) + return 1.0 / (temperature**exp) class ElementWiseFourierEmbed(nn.Module): @@ -258,13 +255,9 @@ class PackedExpertLinear(nn.Module): self.out_features = out_features self.num_experts = num_experts self.use_bias = bias - self.weight = nn.Parameter( - torch.empty(out_features * num_experts, in_features, dtype=dtype) - ) + self.weight = nn.Parameter(torch.empty(out_features * num_experts, in_features, dtype=dtype)) if bias: - self.bias = nn.Parameter( - torch.empty(out_features * num_experts, dtype=dtype) - ) + self.bias = nn.Parameter(torch.empty(out_features * num_experts, dtype=dtype)) else: self.register_parameter("bias", None) @@ -275,15 +268,10 @@ class PackedExpertLinear(nn.Module): ) -> torch.Tensor: if self.num_experts == 1: return F.linear(x, self.weight, self.bias) - assert modality_dispatcher is not None, ( - "PackedExpertLinear with num_experts>1 requires a dispatcher" - ) + assert modality_dispatcher is not None, ("PackedExpertLinear with num_experts>1 requires a dispatcher") parts = modality_dispatcher.dispatch(x) w_chunks = self.weight.chunk(self.num_experts, dim=0) - b_chunks = ( - self.bias.chunk(self.num_experts, dim=0) - if self.bias is not None else [None] * self.num_experts - ) + b_chunks = (self.bias.chunk(self.num_experts, dim=0) if self.bias is not None else [None] * self.num_experts) for i in range(self.num_experts): parts[i] = F.linear(parts[i], w_chunks[i], b_chunks[i]) return modality_dispatcher.undispatch(*parts) @@ -313,18 +301,19 @@ class MagiAttention(nn.Module): super().__init__() self.cfg = cfg self.gating_size = cfg.num_heads_q if cfg.enable_attn_gating else 0 - qkv_out = ( - cfg.num_heads_q * cfg.head_dim - + 2 * cfg.num_heads_kv * cfg.head_dim - + self.gating_size - ) + qkv_out = (cfg.num_heads_q * cfg.head_dim + 2 * cfg.num_heads_kv * cfg.head_dim + self.gating_size) self.pre_norm = MultiModalityRMSNorm(cfg.hidden_size, num_modality=cfg.num_modality) self.linear_qkv = PackedExpertLinear( - cfg.hidden_size, qkv_out, num_experts=cfg.num_modality, bias=False, + cfg.hidden_size, + qkv_out, + num_experts=cfg.num_modality, + bias=False, ) self.linear_proj = PackedExpertLinear( - cfg.num_heads_q * cfg.head_dim, cfg.hidden_size, - num_experts=cfg.num_modality, bias=False, + cfg.num_heads_q * cfg.head_dim, + cfg.hidden_size, + num_experts=cfg.num_modality, + bias=False, ) self.q_norm = MultiModalityRMSNorm(cfg.head_dim, num_modality=cfg.num_modality) self.k_norm = MultiModalityRMSNorm(cfg.head_dim, num_modality=cfg.num_modality) @@ -389,10 +378,8 @@ class MagiAttention(nn.Module): if num_frames <= 0 or num_video_tokens <= 0: return self._sdpa(q, k, v) if num_video_tokens % num_frames != 0: - raise ValueError( - f"MagiHuman local attention expects video tokens divisible by " - f"frames, got {num_video_tokens=} and {num_frames=}." - ) + raise ValueError(f"MagiHuman local attention expects video tokens divisible by " + f"frames, got {num_video_tokens=} and {num_frames=}.") token_per_frame = num_video_tokens // num_frames out = torch.zeros( @@ -447,7 +434,9 @@ class MagiAttention(nn.Module): h = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(orig_dtype) qkv = self.linear_qkv(h, modality_dispatcher=modality_dispatcher).float() q, k, v, g = torch.split( - qkv, [self.q_size, self.kv_size, self.kv_size, self.gating_size], dim=-1, + qkv, + [self.q_size, self.kv_size, self.kv_size, self.gating_size], + dim=-1, ) q = q.view(-1, self.cfg.num_heads_q, self.cfg.head_dim) k = k.view(-1, self.cfg.num_heads_kv, self.cfg.head_dim) @@ -521,23 +510,29 @@ class MagiAttention(nn.Module): class MLPSubConfig: hidden_size: int intermediate_size: int - activation: str # "swiglu7" or "gelu7" + activation: str # "swiglu7" or "gelu7" num_modality: int gated: bool class MagiMLP(nn.Module): + def __init__(self, cfg: MLPSubConfig): super().__init__() self.cfg = cfg self.pre_norm = MultiModalityRMSNorm(cfg.hidden_size, num_modality=cfg.num_modality) up_out = cfg.intermediate_size * 2 if cfg.gated else cfg.intermediate_size self.up_gate_proj = PackedExpertLinear( - cfg.hidden_size, up_out, num_experts=cfg.num_modality, bias=False, + cfg.hidden_size, + up_out, + num_experts=cfg.num_modality, + bias=False, ) self.down_proj = PackedExpertLinear( - cfg.intermediate_size, cfg.hidden_size, - num_experts=cfg.num_modality, bias=False, + cfg.intermediate_size, + cfg.hidden_size, + num_experts=cfg.num_modality, + bias=False, ) self._act = swiglu7 if cfg.activation == "swiglu7" else gelu7 @@ -555,21 +550,23 @@ class MagiMLP(nn.Module): class MagiTransformerLayer(nn.Module): + def __init__(self, arch: MagiHumanArchConfig, layer_idx: int): super().__init__() num_modality = 3 if layer_idx in arch.mm_layers else 1 self.post_norm = layer_idx in arch.post_norm_layers self.layer_idx = layer_idx - self.attention = MagiAttention(AttentionSubConfig( - hidden_size=arch.hidden_size, - num_heads_q=arch.num_attention_heads, - num_heads_kv=arch.num_heads_kv, - head_dim=arch.head_dim, - num_modality=num_modality, - enable_attn_gating=arch.enable_attn_gating, - use_local_attn=layer_idx in arch.local_attn_layers, - )) + self.attention = MagiAttention( + AttentionSubConfig( + hidden_size=arch.hidden_size, + num_heads_q=arch.num_attention_heads, + num_heads_kv=arch.num_heads_kv, + head_dim=arch.head_dim, + num_modality=num_modality, + enable_attn_gating=arch.enable_attn_gating, + use_local_attn=layer_idx in arch.local_attn_layers, + )) is_gelu7 = layer_idx in arch.gelu7_layers if is_gelu7: @@ -581,13 +578,14 @@ class MagiTransformerLayer(nn.Module): gated = True activation = "swiglu7" - self.mlp = MagiMLP(MLPSubConfig( - hidden_size=arch.hidden_size, - intermediate_size=intermediate, - activation=activation, - num_modality=num_modality, - gated=gated, - )) + self.mlp = MagiMLP( + MLPSubConfig( + hidden_size=arch.hidden_size, + intermediate_size=intermediate, + activation=activation, + num_modality=num_modality, + gated=gated, + )) if self.post_norm: self.attn_post_norm = MultiModalityRMSNorm(arch.hidden_size, num_modality=num_modality) @@ -604,7 +602,11 @@ class MagiTransformerLayer(nn.Module): num_frames: int | None = None, ) -> torch.Tensor: attn_out = self.attention( - hidden_states, rope, permute_mapping, inv_permute_mapping, modality_dispatcher, + hidden_states, + rope, + permute_mapping, + inv_permute_mapping, + modality_dispatcher, num_video_tokens=num_video_tokens, num_frames=num_frames, ) @@ -624,17 +626,27 @@ class MagiTransformerLayer(nn.Module): class MagiAdapter(nn.Module): + def __init__(self, arch: MagiHumanArchConfig): super().__init__() # Embedders stay in fp32 to match the reference dtype exactly. self.video_embedder = nn.Linear( - arch.video_in_channels, arch.hidden_size, bias=True, dtype=torch.float32, + arch.video_in_channels, + arch.hidden_size, + bias=True, + dtype=torch.float32, ) self.text_embedder = nn.Linear( - arch.text_in_channels, arch.hidden_size, bias=True, dtype=torch.float32, + arch.text_in_channels, + arch.hidden_size, + bias=True, + dtype=torch.float32, ) self.audio_embedder = nn.Linear( - arch.audio_in_channels, arch.hidden_size, bias=True, dtype=torch.float32, + arch.audio_in_channels, + arch.hidden_size, + bias=True, + dtype=torch.float32, ) self.rope = ElementWiseFourierEmbed(arch.head_dim) # RoPE cache: coords_mapping is the same tensor object across timesteps @@ -666,18 +678,16 @@ class MagiAdapter(nn.Module): a_w = self.audio_embedder.weight t_w = self.text_embedder.weight out = torch.zeros( - x.shape[0], self.video_embedder.out_features, - device=x.device, dtype=v_w.dtype, + x.shape[0], + self.video_embedder.out_features, + device=x.device, + dtype=v_w.dtype, ) - out[text_mask] = self.text_embedder( - x[text_mask, : self.text_embedder.in_features].to(t_w.dtype) - ).to(out.dtype) - out[audio_mask] = self.audio_embedder( - x[audio_mask, : self.audio_embedder.in_features].to(a_w.dtype) - ).to(out.dtype) - out[video_mask] = self.video_embedder( - x[video_mask, : self.video_embedder.in_features].to(v_w.dtype) - ).to(out.dtype) + out[text_mask] = self.text_embedder(x[text_mask, :self.text_embedder.in_features].to(t_w.dtype)).to(out.dtype) + out[audio_mask] = self.audio_embedder(x[audio_mask, :self.audio_embedder.in_features].to(a_w.dtype)).to( + out.dtype) + out[video_mask] = self.video_embedder(x[video_mask, :self.video_embedder.in_features].to(v_w.dtype)).to( + out.dtype) return out, rope @@ -693,9 +703,7 @@ class _TransformerBlock(nn.Module): def __init__(self, arch: MagiHumanArchConfig): super().__init__() - self.layers = nn.ModuleList([ - MagiTransformerLayer(arch, i) for i in range(arch.num_layers) - ]) + self.layers = nn.ModuleList([MagiTransformerLayer(arch, i) for i in range(arch.num_layers)]) def configure_local_attention( self, @@ -770,10 +778,16 @@ class MagiHumanDiT(BaseDiT): self.final_norm_video = MultiModalityRMSNorm(arch.hidden_size) self.final_norm_audio = MultiModalityRMSNorm(arch.hidden_size) self.final_linear_video = nn.Linear( - arch.hidden_size, arch.video_in_channels, bias=False, dtype=torch.float32, + arch.hidden_size, + arch.video_in_channels, + bias=False, + dtype=torch.float32, ) self.final_linear_audio = nn.Linear( - arch.hidden_size, arch.audio_in_channels, bias=False, dtype=torch.float32, + arch.hidden_size, + arch.audio_in_channels, + bias=False, + dtype=torch.float32, ) # Dispatcher + mask cache: modality_mapping is the same tensor object # across all timesteps in the denoising loop; data_ptr()+shape+dtype+device @@ -840,7 +854,8 @@ class MagiHumanDiT(BaseDiT): x = ModalityDispatcher.permute(x, dispatcher.permute_mapping) x = self.block( - x, rope, + x, + rope, permute_mapping=dispatcher.permute_mapping, inv_permute_mapping=dispatcher.inv_permute_mapping, modality_dispatcher=dispatcher, @@ -859,8 +874,8 @@ class MagiHumanDiT(BaseDiT): max_ch = max(self.arch.video_in_channels, self.arch.audio_in_channels) out = torch.zeros(x.shape[0], max_ch, device=x.device, dtype=x.dtype) - out[video_mask, : self.arch.video_in_channels] = x_video.to(out.dtype) - out[audio_mask, : self.arch.audio_in_channels] = x_audio.to(out.dtype) + out[video_mask, :self.arch.video_in_channels] = x_video.to(out.dtype) + out[audio_mask, :self.arch.audio_in_channels] = x_audio.to(out.dtype) return out diff --git a/fastvideo/models/dits/matrixgame2/action_module.py b/fastvideo/models/dits/matrixgame2/action_module.py index b7fc40511..4b2558178 100644 --- a/fastvideo/models/dits/matrixgame2/action_module.py +++ b/fastvideo/models/dits/matrixgame2/action_module.py @@ -18,11 +18,8 @@ from fastvideo.platforms import AttentionBackendEnum from .utils import retain_kv_with_sink - DISABLE_COMPILE = False -flex_attention = torch.compile( - flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs" -) +flex_attention = torch.compile(flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") def _get_nd_rotary_pos_embed_matrixgame2( @@ -54,8 +51,8 @@ def _apply_rotary_emb_qk( seq_len = xq.shape[1] # Slice frequencies based on offset - cos = freqs_cos[start_offset : start_offset + seq_len] # [S, D] - sin = freqs_sin[start_offset : start_offset + seq_len] # [S, D] + cos = freqs_cos[start_offset:start_offset + seq_len] # [S, D] + sin = freqs_sin[start_offset:start_offset + seq_len] # [S, D] # Move to device cos = cos.to(xq.device) @@ -128,29 +125,19 @@ def _update_kv_cache_and_attend( Attention output tensor """ current_start = start_frame - current_end = current_start + ( - k.shape[1] if use_k_for_num_tokens else q.shape[1] - ) + current_end = current_start + (k.shape[1] if use_k_for_num_tokens else q.shape[1]) - assert ( - k.shape[1] if use_k_for_num_tokens else q.shape[1] - ) == num_frame_per_block + assert (k.shape[1] if use_k_for_num_tokens else q.shape[1]) == num_frame_per_block max_attention_size = local_attn_size sink_tokens = sink_size * 1 kv_cache_size = kv_cache["k"].shape[1] num_new_tokens = k.shape[1] if use_k_for_num_tokens else q.shape[1] - original_global_end_index = ( - int(kv_cache["global_end_index"].item()) - if isinstance(kv_cache["global_end_index"], torch.Tensor) - else int(kv_cache["global_end_index"]) - ) - original_local_end_index = ( - int(kv_cache["local_end_index"].item()) - if isinstance(kv_cache["local_end_index"], torch.Tensor) - else int(kv_cache["local_end_index"]) - ) + original_global_end_index = (int(kv_cache["global_end_index"].item()) if isinstance( + kv_cache["global_end_index"], torch.Tensor) else int(kv_cache["global_end_index"])) + original_local_end_index = (int(kv_cache["local_end_index"].item()) if isinstance( + kv_cache["local_end_index"], torch.Tensor) else int(kv_cache["local_end_index"])) if torch.is_grad_enabled(): kv_cache["k"] = kv_cache["k"].detach().clone() @@ -160,60 +147,31 @@ def _update_kv_cache_and_attend( kv_cache["v"] = kv_cache["v"].detach() # Check if we need to evict tokens - if (current_end > original_global_end_index) and ( - num_new_tokens + original_local_end_index > kv_cache_size - ): - num_evicted_tokens = ( - num_new_tokens + original_local_end_index - kv_cache_size - ) - num_rolled_tokens = ( - original_local_end_index - - num_evicted_tokens - - sink_tokens - ) + if (current_end > original_global_end_index) and (num_new_tokens + original_local_end_index > kv_cache_size): + num_evicted_tokens = (num_new_tokens + original_local_end_index - kv_cache_size) + num_rolled_tokens = (original_local_end_index - num_evicted_tokens - sink_tokens) # Roll k cache - kv_cache["k"][:, sink_tokens : sink_tokens + num_rolled_tokens] = ( - kv_cache["k"][ - :, - sink_tokens + num_evicted_tokens : sink_tokens - + num_evicted_tokens - + num_rolled_tokens, - ].clone() - ) + kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = (kv_cache["k"][ + :, + sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens, + ].clone()) # Roll v cache - kv_cache["v"][:, sink_tokens : sink_tokens + num_rolled_tokens] = ( - kv_cache["v"][ - :, - sink_tokens + num_evicted_tokens : sink_tokens - + num_evicted_tokens - + num_rolled_tokens, - ].clone() - ) + kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = (kv_cache["v"][ + :, + sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens, + ].clone()) # Calculate indices with eviction adjustment - local_end_index = ( - original_local_end_index - + current_end - - original_global_end_index - - num_evicted_tokens - ) + local_end_index = (original_local_end_index + current_end - original_global_end_index - num_evicted_tokens) local_start_index = local_end_index - num_new_tokens else: # Calculate indices without eviction - local_end_index = ( - original_local_end_index - + current_end - - original_global_end_index - ) + local_end_index = (original_local_end_index + current_end - original_global_end_index) local_start_index = local_end_index - num_new_tokens # Store new k, v in cache if store_first_only: - kv_cache["k"][:, local_start_index:local_end_index] = k[ - : kv_cache["k"].shape[0] - ] - kv_cache["v"][:, local_start_index:local_end_index] = v[ - : kv_cache["v"].shape[0] - ] + kv_cache["k"][:, local_start_index:local_end_index] = k[:kv_cache["k"].shape[0]] + kv_cache["v"][:, local_start_index:local_end_index] = v[:kv_cache["v"].shape[0]] else: kv_cache["k"][:, local_start_index:local_end_index] = k kv_cache["v"][:, local_start_index:local_end_index] = v @@ -275,12 +233,8 @@ class ActionModule(nn.Module): super().__init__() # Initialize mutable defaults patch_size = patch_size if patch_size is not None else [1, 2, 2] - rope_dim_list = ( - rope_dim_list if rope_dim_list is not None else [8, 28, 28] - ) - mouse_qk_dim_list = ( - mouse_qk_dim_list if mouse_qk_dim_list is not None else [8, 28, 28] - ) + rope_dim_list = (rope_dim_list if rope_dim_list is not None else [8, 28, 28]) + mouse_qk_dim_list = (mouse_qk_dim_list if mouse_qk_dim_list is not None else [8, 28, 28]) blocks = blocks if blocks is not None else [] self.local_attn_size = local_attn_size self.sink_size = sink_size @@ -302,8 +256,7 @@ class ActionModule(nn.Module): c = mouse_hidden_dim self.mouse_mlp = nn.Sequential( nn.Linear( - mouse_dim_in * vae_time_compression_ratio * windows_size - + img_hidden_size, + mouse_dim_in * vae_time_compression_ratio * windows_size + img_hidden_size, c, bias=True, ), @@ -314,64 +267,42 @@ class ActionModule(nn.Module): head_dim = c // heads_num self.t_qkv = ReplicatedLinear(c, c * 3, bias=qkv_bias) - self.img_attn_q_norm = ( - RMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity() - ) - self.img_attn_k_norm = ( - RMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity() - ) - self.proj_mouse = ReplicatedLinear( - c, img_hidden_size, bias=qkv_bias - ) + self.img_attn_q_norm = (RMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity()) + self.img_attn_k_norm = (RMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity()) + self.proj_mouse = ReplicatedLinear(c, img_hidden_size, bias=qkv_bias) if self.enable_keyboard: head_dim_key = keyboard_hidden_dim // heads_num - self.key_attn_q_norm = ( - RMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity() - ) - self.key_attn_k_norm = ( - RMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity() - ) + self.key_attn_q_norm = (RMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity()) + self.key_attn_k_norm = (RMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity()) - self.mouse_attn_q = ReplicatedLinear( - img_hidden_size, keyboard_hidden_dim, bias=qkv_bias - ) + self.mouse_attn_q = ReplicatedLinear(img_hidden_size, keyboard_hidden_dim, bias=qkv_bias) self.keyboard_attn_kv = ReplicatedLinear( hidden_size * windows_size * vae_time_compression_ratio, keyboard_hidden_dim * 2, bias=qkv_bias, ) - self.proj_keyboard = ReplicatedLinear( - keyboard_hidden_dim, img_hidden_size, bias=qkv_bias - ) + self.proj_keyboard = ReplicatedLinear(keyboard_hidden_dim, img_hidden_size, bias=qkv_bias) - self.mouse_attn_layer = ( - LocalAttention( - num_heads=heads_num, - head_size=mouse_hidden_dim // heads_num, - causal=False, - supported_attention_backends=( - AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA, - ), - ) - if self.enable_mouse - else None - ) + self.mouse_attn_layer = (LocalAttention( + num_heads=heads_num, + head_size=mouse_hidden_dim // heads_num, + causal=False, + supported_attention_backends=( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ), + ) if self.enable_mouse else None) - self.keyboard_attn_layer = ( - LocalAttention( - num_heads=heads_num, - head_size=keyboard_hidden_dim // heads_num, - causal=False, - supported_attention_backends=( - AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA, - ), - ) - if self.enable_keyboard - else None - ) + self.keyboard_attn_layer = (LocalAttention( + num_heads=heads_num, + head_size=keyboard_hidden_dim // heads_num, + causal=False, + supported_attention_backends=( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ), + ) if self.enable_keyboard else None) self.vae_time_compression_ratio = vae_time_compression_ratio self.windows_size = windows_size @@ -424,33 +355,20 @@ class ActionModule(nn.Module): if isinstance(self.patch_size, int): assert all(s % self.patch_size == 0 for s in latents_size), ( f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), " - f"but got {latents_size}." - ) + f"but got {latents_size}.") rope_sizes = [s // self.patch_size for s in latents_size] elif isinstance(self.patch_size, list): - assert all( - s % self.patch_size[idx] == 0 - for idx, s in enumerate(latents_size) - ), ( + assert all(s % self.patch_size[idx] == 0 for idx, s in enumerate(latents_size)), ( f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), " - f"but got {latents_size}." - ) - rope_sizes = [ - s // self.patch_size[idx] for idx, s in enumerate(latents_size) - ] + f"but got {latents_size}.") + rope_sizes = [s // self.patch_size[idx] for idx, s in enumerate(latents_size)] if len(rope_sizes) != target_ndim: - rope_sizes = [1] * ( - target_ndim - len(rope_sizes) - ) + rope_sizes # time axis + rope_sizes = [1] * (target_ndim - len(rope_sizes)) + rope_sizes # time axis if rope_dim_list is None: - rope_dim_list = [ - head_dim // target_ndim for _ in range(target_ndim) - ] - assert sum(rope_dim_list) == head_dim, ( - "sum(rope_dim_list) should equal to head_dim of attention layer" - ) + rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] + assert sum(rope_dim_list) == head_dim, ("sum(rope_dim_list) should equal to head_dim of attention layer") # Use Matrix-Game wrapper for FastVideo's function freqs_cos, freqs_sin = _get_nd_rotary_pos_embed_matrixgame2( rope_dim_list, @@ -458,17 +376,9 @@ class ActionModule(nn.Module): theta=self.rope_theta, theta_rescale_factor=1, ) - return freqs_cos[ - -video_length - * rope_sizes[1] - * rope_sizes[2] - // self.patch_size[0] : - ], freqs_sin[ - -video_length - * rope_sizes[1] - * rope_sizes[2] - // self.patch_size[0] : - ] + return freqs_cos[-video_length * rope_sizes[1] * rope_sizes[2] // + self.patch_size[0]:], freqs_sin[-video_length * rope_sizes[1] * rope_sizes[2] // + self.patch_size[0]:] def _forward_mouse( self, @@ -494,67 +404,51 @@ class ActionModule(nn.Module): if is_causal and kv_cache_mouse is not None: mouse_condition = mouse_condition[ :, - self.vae_time_compression_ratio - * (N_feats - num_frame_per_block - self.windows_size) - + pad_t :, + self.vae_time_compression_ratio * (N_feats - num_frame_per_block - self.windows_size) + pad_t:, :, ] group_mouse = [ mouse_condition[ :, - self.vae_time_compression_ratio * (i - self.windows_size) - + pad_t : i * self.vae_time_compression_ratio + pad_t, + self.vae_time_compression_ratio * (i - self.windows_size) + + pad_t:i * self.vae_time_compression_ratio + pad_t, :, - ] - for i in range(num_frame_per_block) + ] for i in range(num_frame_per_block) ] else: local_num_frames = tt group_mouse = [ mouse_condition[ :, - self.vae_time_compression_ratio * (i - self.windows_size) - + pad_t : i * self.vae_time_compression_ratio + pad_t, + self.vae_time_compression_ratio * (i - self.windows_size) + + pad_t:i * self.vae_time_compression_ratio + pad_t, :, - ] - for i in range(local_num_frames) + ] for i in range(local_num_frames) ] group_mouse = torch.stack(group_mouse, dim=1) - actual_num_frames = group_mouse.shape[ - 1 - ] # Use actual stacked frame count + actual_num_frames = group_mouse.shape[1] # Use actual stacked frame count S = th * tw - group_mouse = group_mouse.unsqueeze(-1).expand( - B, actual_num_frames, pad_t, C, S - ) - group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape( - B * S, actual_num_frames, pad_t * C - ) + group_mouse = group_mouse.unsqueeze(-1).expand(B, actual_num_frames, pad_t, C, S) + group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape(B * S, actual_num_frames, pad_t * C) group_mouse = torch.cat([hidden_states, group_mouse], dim=-1) group_mouse = self.mouse_mlp(group_mouse) # qkv mouse_qkv, _ = self.t_qkv(group_mouse) - q, k, v = rearrange( - mouse_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num - ) # BHW F H C + q, k, v = rearrange(mouse_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) # BHW F H C q = self.img_attn_q_norm(q).to(v) k = self.img_attn_k_norm(k).to(v) # rope embd # freqs_cis = (self.freqs_cos, self.freqs_sin) - q, k = _apply_rotary_emb_qk( - q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame - ) + q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame) ## TODO: adding cache here if is_causal: if kv_cache_mouse is None: - assert ( - q.shape[0] == k.shape[0] and q.shape[0] % S == 0 - ) # == 880, f"{q.shape[0]},{k.shape[0]}" + assert (q.shape[0] == k.shape[0] and q.shape[0] % S == 0) # == 880, f"{q.shape[0]},{k.shape[0]}" padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1] padded_q = _padding_q_k_v(q, padded_length) padded_k = _padding_q_k_v(k, padded_length) @@ -610,20 +504,17 @@ class ActionModule(nn.Module): if is_causal and kv_cache_keyboard is not None: keyboard_condition = keyboard_condition[ :, - self.vae_time_compression_ratio - * (N_feats - num_frame_per_block - self.windows_size) - + pad_t :, + self.vae_time_compression_ratio * (N_feats - num_frame_per_block - self.windows_size) + pad_t:, :, ] # keyboard_condition[:, self.vae_time_compression_ratio*(start_frame - self.windows_size) + pad_t:start_frame * self.vae_time_compression_ratio + pad_t,:] keyboard_condition = self.keyboard_embed(keyboard_condition) group_keyboard = [ keyboard_condition[ :, - self.vae_time_compression_ratio * (i - self.windows_size) - + pad_t : i * self.vae_time_compression_ratio + pad_t, + self.vae_time_compression_ratio * (i - self.windows_size) + + pad_t:i * self.vae_time_compression_ratio + pad_t, :, - ] - for i in range(num_frame_per_block) + ] for i in range(num_frame_per_block) ] else: keyboard_condition = self.keyboard_embed(keyboard_condition) @@ -631,16 +522,13 @@ class ActionModule(nn.Module): group_keyboard = [ keyboard_condition[ :, - self.vae_time_compression_ratio * (i - self.windows_size) - + pad_t : i * self.vae_time_compression_ratio + pad_t, + self.vae_time_compression_ratio * (i - self.windows_size) + + pad_t:i * self.vae_time_compression_ratio + pad_t, :, - ] - for i in range(local_num_frames) + ] for i in range(local_num_frames) ] group_keyboard = torch.stack(group_keyboard, dim=1) # B F RW C - group_keyboard = group_keyboard.reshape( - shape=(group_keyboard.shape[0], group_keyboard.shape[1], -1) - ) + group_keyboard = group_keyboard.reshape(shape=(group_keyboard.shape[0], group_keyboard.shape[1], -1)) # apply cross attn mouse_q, _ = self.mouse_attn_q(hidden_states) keyboard_kv, _ = self.keyboard_attn_kv(group_keyboard) @@ -650,9 +538,7 @@ class ActionModule(nn.Module): q = mouse_q.view(B, L, self.heads_num, D) B, L, KHD = keyboard_kv.shape - k, v = keyboard_kv.view(B, L, 2, self.heads_num, D).permute( - 2, 0, 1, 3, 4 - ) + k, v = keyboard_kv.view(B, L, 2, self.heads_num, D).permute(2, 0, 1, 3, 4) # Compute cu_squlens and max_seqlen for flash attention # qk norm @@ -666,9 +552,7 @@ class ActionModule(nn.Module): B, TS, H, D = q.shape T_ = TS // S q = q.view(B, T_, S, H, D).transpose(1, 2).reshape(B * S, T_, H, D) - q, k = _apply_rotary_emb_qk( - q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame - ) + q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame) k = k.repeat_interleave(S, dim=0) v = v.repeat_interleave(S, dim=0) @@ -688,9 +572,7 @@ class ActionModule(nn.Module): block_mask=block_mask_keyboard, )[:, :, :-padded_length].transpose(2, 1) else: - assert ( - k.shape[0] == S - ) # BS == 1 or the cache should not be saved/ load method should be modified + assert (k.shape[0] == S) # BS == 1 or the cache should not be saved/ load method should be modified attn = _update_kv_cache_and_attend( q, k, @@ -767,21 +649,15 @@ class ActionModule(nn.Module): target_device = x.device target_dtype = x.dtype if mouse_condition is not None: - mouse_condition = mouse_condition.to( - device=target_device, dtype=target_dtype - ) + mouse_condition = mouse_condition.to(device=target_device, dtype=target_dtype) if keyboard_condition is not None: - keyboard_condition = keyboard_condition.to( - device=target_device, dtype=target_dtype - ) + keyboard_condition = keyboard_condition.to(device=target_device, dtype=target_dtype) else: return x B, N_frames, C = keyboard_condition.shape assert tt * th * tw == x.shape[1] - assert ( - (N_frames - 1) + self.vae_time_compression_ratio - ) % self.vae_time_compression_ratio == 0 + assert ((N_frames - 1) + self.vae_time_compression_ratio) % self.vae_time_compression_ratio == 0 N_feats = int((N_frames - 1) / self.vae_time_compression_ratio) + 1 # Lazy initialization of freqs on first forward pass @@ -799,16 +675,13 @@ class ActionModule(nn.Module): freqs_cis = (self._freqs_cos, self._freqs_sin) if is_causal: - assert (N_feats == tt and kv_cache_mouse is None) or ( - (N_frames - 1) // self.vae_time_compression_ratio + 1 - == start_frame + num_frame_per_block - ) + assert (N_feats == tt and kv_cache_mouse is None) or ((N_frames - 1) // self.vae_time_compression_ratio + 1 + == start_frame + num_frame_per_block) # For non-causal (training), we trust that the caller provides correctly shaped inputs if self.enable_mouse and mouse_condition is not None: - hidden_states = rearrange( - x, "B (T S) C -> (B S) T C", T=tt, S=th * tw - ) # 65*272*480 -> 17*(272//16)*(480//16) -> 8670 + hidden_states = rearrange(x, "B (T S) C -> (B S) T C", T=tt, + S=th * tw) # 65*272*480 -> 17*(272//16)*(480//16) -> 8670 B, N_frames, C = mouse_condition.shape else: hidden_states = x diff --git a/fastvideo/models/dits/matrixgame2/causal_model.py b/fastvideo/models/dits/matrixgame2/causal_model.py index 1b98e2910..ba193dc18 100644 --- a/fastvideo/models/dits/matrixgame2/causal_model.py +++ b/fastvideo/models/dits/matrixgame2/causal_model.py @@ -13,9 +13,7 @@ from torch.nn.attention.flex_attention import create_block_mask, flex_attention from torch.nn.attention.flex_attention import BlockMask from fastvideo.forward_context import get_forward_context -flex_attention = torch.compile( - flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs" -) +flex_attention = torch.compile(flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") from fastvideo.attention import LocalAttention from fastvideo.configs.models.dits.matrixgame2 import MatrixGame2WanVideoConfig @@ -51,9 +49,7 @@ from .utils import retain_kv_with_sink logger = init_logger(__name__) -def causal_rope_apply( - x: torch.Tensor, grid_sizes: tuple[int, int, int], freqs: torch.Tensor, start_frame: int = 0 -): +def causal_rope_apply(x: torch.Tensor, grid_sizes: tuple[int, int, int], freqs: torch.Tensor, start_frame: int = 0): n, c = x.size(2), x.size(3) // 2 # split freqs @@ -67,14 +63,10 @@ def causal_rope_apply( seq_len = f * h * w # precompute multipliers - x_i = torch.view_as_complex( - x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2) - ) + x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)) freqs_i = torch.cat( [ - freqs[0][start_frame : start_frame + f] - .view(f, 1, 1, -1) - .expand(f, h, w, -1), + freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), ], @@ -94,14 +86,15 @@ def rope_params(max_seq_len, dim, theta=10000): assert dim % 2 == 0 freqs = torch.outer( torch.arange(max_seq_len), - 1.0 - / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)), + 1.0 / torch.pow(theta, + torch.arange(0, dim, 2).to(torch.float64).div(dim)), ) freqs = torch.polar(torch.ones_like(freqs), freqs) return freqs class CausalMatrixGame2TimeImageEmbedding(nn.Module): + def __init__( self, dim: int, @@ -110,12 +103,8 @@ class CausalMatrixGame2TimeImageEmbedding(nn.Module): ): super().__init__() - self.time_embedder = TimestepEmbedder( - dim, frequency_embedding_size=time_freq_dim, act_layer="silu" - ) - self.time_modulation = ModulateProjection( - dim, factor=6, act_layer="silu" - ) + self.time_embedder = TimestepEmbedder(dim, frequency_embedding_size=time_freq_dim, act_layer="silu") + self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") self.image_embedder = None if image_embed_dim is not None: @@ -133,14 +122,13 @@ class CausalMatrixGame2TimeImageEmbedding(nn.Module): if encoder_hidden_states_image is not None: assert self.image_embedder is not None - encoder_hidden_states_image = self.image_embedder( - encoder_hidden_states_image - ) + encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image) return temb, timestep_proj, None, encoder_hidden_states_image class CausalMatrixGame2SelfAttention(nn.Module): + def __init__( self, dim: int, @@ -213,12 +201,8 @@ class CausalMatrixGame2SelfAttention(nn.Module): # relativistic defers roping until the cache window is known (and caches raw k) relativistic = self.rope_cache_policy == "relativistic" and kv_cache is not None if not relativistic: - roped_query = causal_rope_apply( - q, grid_sizes, freqs, start_frame=start_frame - ).type_as(v) - roped_key = causal_rope_apply( - k, grid_sizes, freqs, start_frame=start_frame - ).type_as(v) + roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=start_frame).type_as(v) + roped_key = causal_rope_apply(k, grid_sizes, freqs, start_frame=start_frame).type_as(v) if kv_cache is None: padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] @@ -271,19 +255,14 @@ class CausalMatrixGame2SelfAttention(nn.Module): else: # Fallback: assume q.shape[1] is single frame (shouldn't happen in causal mode) frame_seqlen = q.shape[1] - logger.warning( - "grid_sizes not provided, using q.shape[1] as frame_seqlen" - ) + logger.warning("grid_sizes not provided, using q.shape[1] as frame_seqlen") current_end = current_start + q.shape[1] sink_tokens = self.sink_size * frame_seqlen # Compute max_attention_size dynamically based on actual frame_seqlen - max_attention_size = ( - 15 * frame_seqlen - if self.local_attn_size == -1 - else self.local_attn_size * frame_seqlen - ) + max_attention_size = (15 * frame_seqlen if self.local_attn_size == -1 else self.local_attn_size * + frame_seqlen) if torch.is_grad_enabled(): kv_cache["k"] = kv_cache["k"].detach().clone() @@ -295,53 +274,28 @@ class CausalMatrixGame2SelfAttention(nn.Module): kv_cache_size = kv_cache["k"].shape[1] num_new_tokens = q.shape[1] stored_key = k if relativistic else roped_key # raw vs roped in cache - global_end_index = ( - int(kv_cache["global_end_index"].item()) - if isinstance(kv_cache["global_end_index"], torch.Tensor) - else int(kv_cache["global_end_index"]) - ) - local_end_index_prev = ( - int(kv_cache["local_end_index"].item()) - if isinstance(kv_cache["local_end_index"], torch.Tensor) - else int(kv_cache["local_end_index"]) - ) + global_end_index = (int(kv_cache["global_end_index"].item()) if isinstance( + kv_cache["global_end_index"], torch.Tensor) else int(kv_cache["global_end_index"])) + local_end_index_prev = (int(kv_cache["local_end_index"].item()) if isinstance( + kv_cache["local_end_index"], torch.Tensor) else int(kv_cache["local_end_index"])) - if (current_end > global_end_index) and ( - num_new_tokens + local_end_index_prev > kv_cache_size - ): - num_evicted_tokens = ( - num_new_tokens + local_end_index_prev - kv_cache_size - ) - num_rolled_tokens = ( - local_end_index_prev - num_evicted_tokens - sink_tokens - ) - kv_cache["k"][ - :, sink_tokens : sink_tokens + num_rolled_tokens - ] = kv_cache["k"][ + if (current_end > global_end_index) and (num_new_tokens + local_end_index_prev > kv_cache_size): + num_evicted_tokens = (num_new_tokens + local_end_index_prev - kv_cache_size) + num_rolled_tokens = (local_end_index_prev - num_evicted_tokens - sink_tokens) + kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = kv_cache["k"][ :, - sink_tokens + num_evicted_tokens : sink_tokens - + num_evicted_tokens - + num_rolled_tokens, + sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens, ].clone() - kv_cache["v"][ - :, sink_tokens : sink_tokens + num_rolled_tokens - ] = kv_cache["v"][ + kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = kv_cache["v"][ :, - sink_tokens + num_evicted_tokens : sink_tokens - + num_evicted_tokens - + num_rolled_tokens, + sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens, ].clone() - local_end_index = ( - local_end_index_prev - + current_end - global_end_index - num_evicted_tokens - ) + local_end_index = (local_end_index_prev + current_end - global_end_index - num_evicted_tokens) local_start_index = local_end_index - num_new_tokens kv_cache["k"][:, local_start_index:local_end_index] = stored_key kv_cache["v"][:, local_start_index:local_end_index] = v else: - local_end_index = ( - local_end_index_prev + current_end - global_end_index - ) + local_end_index = (local_end_index_prev + current_end - global_end_index) local_start_index = local_end_index - num_new_tokens kv_cache["k"][:, local_start_index:local_end_index] = stored_key kv_cache["v"][:, local_start_index:local_end_index] = v @@ -354,15 +308,12 @@ class CausalMatrixGame2SelfAttention(nn.Module): ) if relativistic: - window_len, query_lo, _ = relativistic_window_offsets( - local_end_index, num_new_tokens, max_attention_size) + window_len, query_lo, _ = relativistic_window_offsets(local_end_index, num_new_tokens, + max_attention_size) h, w = int(grid_sizes[1]), int(grid_sizes[2]) window_grid = (window_len // frame_seqlen, h, w) - roped_query = causal_rope_apply( - q, grid_sizes, freqs, - start_frame=query_lo // frame_seqlen).type_as(v) - k_for_attn = causal_rope_apply( - k_for_attn, window_grid, freqs, start_frame=0).type_as(v) + roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=query_lo // frame_seqlen).type_as(v) + k_for_attn = causal_rope_apply(k_for_attn, window_grid, freqs, start_frame=0).type_as(v) x = torch.nn.functional.scaled_dot_product_attention( roped_query.transpose(1, 2), @@ -383,6 +334,7 @@ class CausalMatrixGame2SelfAttention(nn.Module): class CausalMatrixGame2TransformerBlock(nn.Module): + def __init__( self, dim: int, @@ -443,13 +395,9 @@ class CausalMatrixGame2TransformerBlock(nn.Module): ) if added_kv_proj_dim is not None: - self.attn2 = MatrixGame2CrossAttention( - dim, num_heads, qk_norm=qk_norm, eps=eps - ) + self.attn2 = MatrixGame2CrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) else: - self.attn2 = WanT2VCrossAttention( - dim, num_heads, qk_norm=qk_norm, eps=eps - ) + self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, norm_type="layer", @@ -458,9 +406,7 @@ class CausalMatrixGame2TransformerBlock(nn.Module): dtype=torch.float32, ) - if len(action_config) > 0 and block_idx in action_config.get( - "blocks", [] - ): + if len(action_config) > 0 and block_idx in action_config.get("blocks", []): self.action_model = ActionModule( heads_num=action_config["heads_num"], hidden_size=action_config["hidden_size"], @@ -474,17 +420,13 @@ class CausalMatrixGame2TransformerBlock(nn.Module): windows_size=action_config["windows_size"], rope_theta=action_config["rope_theta"], rope_dim_list=action_config["rope_dim_list"], - mouse_qk_dim_list=action_config.get( - "mouse_qk_dim_list", [8, 28, 28] - ), + mouse_qk_dim_list=action_config.get("mouse_qk_dim_list", [8, 28, 28]), patch_size=action_config["patch_size"], local_attn_size=local_attn_size, sink_size=sink_size, qk_norm=action_config["qk_norm"], qkv_bias=action_config["qkv_bias"], - vae_time_compression_ratio=action_config[ - "vae_time_compression_ratio" - ], + vae_time_compression_ratio=action_config["vae_time_compression_ratio"], ) else: self.action_model = None @@ -527,17 +469,10 @@ class CausalMatrixGame2TransformerBlock(nn.Module): e = self.scale_shift_table + temb assert e.shape == (bs, num_frames, 6, self.hidden_dim) - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - e.chunk(6, dim=2) - ) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (e.chunk(6, dim=2)) - norm_hidden_states = ( - self.norm1(hidden_states).unflatten( - dim=1, sizes=(num_frames, frame_seqlen) - ) - * (1 + scale_msa) - + shift_msa - ).flatten(1, 2) + norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + (1 + scale_msa) + shift_msa).flatten(1, 2) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -567,9 +502,8 @@ class CausalMatrixGame2TransformerBlock(nn.Module): attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.tensor([0], device=hidden_states.device) - norm_hidden_states, hidden_states = self.self_attn_residual_norm( - hidden_states, attn_output, gate_msa, null_shift, null_scale - ) + norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, + null_shift, null_scale) # norm3 weights are loaded into self_attn_residual_norm.norm attn_output = self.attn2( @@ -602,15 +536,8 @@ class CausalMatrixGame2TransformerBlock(nn.Module): num_frame_per_block=num_frame_per_block, ) - ff_output = self.ffn( - ( - self.norm1(hidden_states).unflatten( - dim=1, sizes=(num_frames, frame_seqlen) - ) - * (1 + c_scale_msa) - + c_shift_msa - ).flatten(1, 2) - ) + ff_output = self.ffn((self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + (1 + c_scale_msa) + c_shift_msa).flatten(1, 2)) hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) # Convert back to original dtype before return @@ -625,16 +552,10 @@ class CausalMatrixGame2WanModel(BaseDiT): _fsdp_shard_conditions = _DEFAULT_MATRIXGAME2_CONFIG._fsdp_shard_conditions _compile_conditions = _DEFAULT_MATRIXGAME2_CONFIG._compile_conditions - _supported_attention_backends = ( - _DEFAULT_MATRIXGAME2_CONFIG._supported_attention_backends - ) + _supported_attention_backends = (_DEFAULT_MATRIXGAME2_CONFIG._supported_attention_backends) param_names_mapping = _DEFAULT_MATRIXGAME2_CONFIG.param_names_mapping - reverse_param_names_mapping = ( - _DEFAULT_MATRIXGAME2_CONFIG.reverse_param_names_mapping - ) - lora_param_names_mapping = ( - _DEFAULT_MATRIXGAME2_CONFIG.lora_param_names_mapping - ) + reverse_param_names_mapping = (_DEFAULT_MATRIXGAME2_CONFIG.reverse_param_names_mapping) + lora_param_names_mapping = (_DEFAULT_MATRIXGAME2_CONFIG.lora_param_names_mapping) def __init__( self, @@ -656,32 +577,20 @@ class CausalMatrixGame2WanModel(BaseDiT): # Long tuning controls the causal window from YAML; consume it here # like Wan so train wrappers do not need runtime patching. arch_cfg = getattr(config, "arch_config", None) - self.local_attn_size = ( - getattr( - arch_cfg, - "local_attn_size", - getattr(config, "local_attn_size", -1), - ) - if arch_cfg - else getattr(config, "local_attn_size", -1) - ) - self.sink_size = ( - getattr(arch_cfg, "sink_size", getattr(config, "sink_size", 0)) - if arch_cfg - else getattr(config, "sink_size", 0) - ) + self.local_attn_size = (getattr( + arch_cfg, + "local_attn_size", + getattr(config, "local_attn_size", -1), + ) if arch_cfg else getattr(config, "local_attn_size", -1)) + self.sink_size = (getattr(arch_cfg, "sink_size", getattr(config, "sink_size", 0)) if arch_cfg else getattr( + config, "sink_size", 0)) if self.sink_size < 0: raise ValueError("sink_size must be non-negative") if self.local_attn_size != -1 and self.sink_size >= self.local_attn_size: - raise ValueError( - "sink_size must be smaller than local_attn_size for " - "MatrixGame2 causal attention") - self.rope_cache_policy = ( - getattr(arch_cfg, "rope_cache_policy", - getattr(config, "rope_cache_policy", "absolute")) - if arch_cfg - else getattr(config, "rope_cache_policy", "absolute") - ) + raise ValueError("sink_size must be smaller than local_attn_size for " + "MatrixGame2 causal attention") + self.rope_cache_policy = (getattr(arch_cfg, "rope_cache_policy", getattr( + config, "rope_cache_policy", "absolute")) if arch_cfg else getattr(config, "rope_cache_policy", "absolute")) # 1. Patch & position embedding self.patch_embedding = PatchEmbed( @@ -700,34 +609,27 @@ class CausalMatrixGame2WanModel(BaseDiT): # 2.1. Get action config arch_cfg = getattr(config, "arch_config", None) - self.action_config = ( - getattr(arch_cfg, "action_config", {}) - if arch_cfg is not None - else {} - ) + self.action_config = (getattr(arch_cfg, "action_config", {}) if arch_cfg is not None else {}) # 3. Transformer blocks - self.blocks = nn.ModuleList( - [ - CausalMatrixGame2TransformerBlock( - inner_dim, - config.ffn_dim, - config.num_attention_heads, - self.local_attn_size, - self.sink_size, - config.qk_norm, - config.cross_attn_norm, - config.eps, - config.added_kv_proj_dim, - self._supported_attention_backends, - prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", - action_config=self.action_config, - block_idx=i, - rope_cache_policy=self.rope_cache_policy, - ) - for i in range(config.num_layers) - ] - ) + self.blocks = nn.ModuleList([ + CausalMatrixGame2TransformerBlock( + inner_dim, + config.ffn_dim, + config.num_attention_heads, + self.local_attn_size, + self.sink_size, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + self._supported_attention_backends, + prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", + action_config=self.action_config, + block_idx=i, + rope_cache_policy=self.rope_cache_policy, + ) for i in range(config.num_layers) + ]) # 4. Output norm & projection self.norm_out = LayerNormScaleShift( @@ -737,12 +639,8 @@ class CausalMatrixGame2WanModel(BaseDiT): elementwise_affine=False, dtype=torch.float32, ) - self.proj_out = nn.Linear( - inner_dim, config.out_channels * math.prod(config.patch_size) - ) - self.scale_shift_table = nn.Parameter( - torch.randn(1, 2, inner_dim) / inner_dim**0.5 - ) + self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) self.gradient_checkpointing = False @@ -751,15 +649,11 @@ class CausalMatrixGame2WanModel(BaseDiT): self.block_mask_mouse = None self.use_rope_keyboard = True - self.num_frame_per_block = ( - getattr( - arch_cfg, - "num_frames_per_block", - getattr(config, "num_frame_per_block", 1), - ) - if arch_cfg - else getattr(config, "num_frame_per_block", 1) - ) + self.num_frame_per_block = (getattr( + arch_cfg, + "num_frames_per_block", + getattr(config, "num_frame_per_block", 1), + ) if arch_cfg else getattr(config, "num_frame_per_block", 1)) self.__post_init__() @@ -769,16 +663,13 @@ class CausalMatrixGame2WanModel(BaseDiT): local_attn_size: int | None = None, sink_size: int | None = None, ) -> None: - local_attn_size = ( - self.local_attn_size if local_attn_size is None else int(local_attn_size) - ) + local_attn_size = (self.local_attn_size if local_attn_size is None else int(local_attn_size)) sink_size = self.sink_size if sink_size is None else int(sink_size) if sink_size < 0: raise ValueError("sink_size must be non-negative") if local_attn_size != -1 and sink_size >= local_attn_size: - raise ValueError( - "sink_size must be smaller than local_attn_size for " - "MatrixGame2 causal attention") + raise ValueError("sink_size must be smaller than local_attn_size for " + "MatrixGame2 causal attention") self.local_attn_size = local_attn_size self.sink_size = sink_size @@ -821,9 +712,7 @@ class CausalMatrixGame2WanModel(BaseDiT): total_length = num_frames * frame_seqlen padded_length = math.ceil(total_length / 128) * 128 - total_length - ends = torch.zeros( - total_length + padded_length, device=device, dtype=torch.long - ) + ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) frame_indices = torch.arange( start=0, @@ -833,21 +722,15 @@ class CausalMatrixGame2WanModel(BaseDiT): ) for tmp in frame_indices: - ends[tmp : tmp + frame_seqlen * num_frame_per_block] = ( - tmp + frame_seqlen * num_frame_per_block - ) + ends[tmp:tmp + frame_seqlen * num_frame_per_block] = (tmp + frame_seqlen * num_frame_per_block) def attention_mask(b, h, q_idx, kv_idx): if local_attn_size == -1: return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) else: - return ( - (kv_idx < ends[q_idx]) - & ( - (kv_idx < sink_size * frame_seqlen) - | (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen)) - ) - ) | (q_idx == kv_idx) + return ((kv_idx < ends[q_idx]) + & ((kv_idx < sink_size * frame_seqlen) + | (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen)))) | (q_idx == kv_idx) block_mask = create_block_mask( attention_mask, @@ -879,9 +762,7 @@ class CausalMatrixGame2WanModel(BaseDiT): total_length2 = num_frames * frame_seqlen padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2 padded_length_kv2 = math.ceil(num_frames / 32) * 32 - num_frames - ends2 = torch.zeros( - total_length2 + padded_length2, device=device, dtype=torch.long - ) + ends2 = torch.zeros(total_length2 + padded_length2, device=device, dtype=torch.long) frame_indices2 = torch.arange( start=0, @@ -891,17 +772,15 @@ class CausalMatrixGame2WanModel(BaseDiT): ) cnt = num_frame_per_block for tmp in frame_indices2: - ends2[tmp : tmp + frame_seqlen * num_frame_per_block] = cnt + ends2[tmp:tmp + frame_seqlen * num_frame_per_block] = cnt cnt += num_frame_per_block def attention_mask2(b, h, q_idx, kv_idx): if local_attn_size == -1: return (kv_idx < ends2[q_idx]) | (q_idx == kv_idx) else: - return ( - (kv_idx < ends2[q_idx]) - & ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size))) - ) | (q_idx == kv_idx) + return ((kv_idx < ends2[q_idx]) + & ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size)))) | (q_idx == kv_idx) block_mask2 = create_block_mask( attention_mask2, @@ -933,9 +812,7 @@ class CausalMatrixGame2WanModel(BaseDiT): total_length2 = num_frames * frame_seqlen padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2 padded_length_kv2 = math.ceil(num_frames / 32) * 32 - num_frames - ends2 = torch.zeros( - total_length2 + padded_length2, device=device, dtype=torch.long - ) + ends2 = torch.zeros(total_length2 + padded_length2, device=device, dtype=torch.long) frame_indices2 = torch.arange( start=0, @@ -945,17 +822,15 @@ class CausalMatrixGame2WanModel(BaseDiT): ) cnt = num_frame_per_block for tmp in frame_indices2: - ends2[tmp : tmp + frame_seqlen * num_frame_per_block] = cnt + ends2[tmp:tmp + frame_seqlen * num_frame_per_block] = cnt cnt += num_frame_per_block def attention_mask2(b, h, q_idx, kv_idx): if local_attn_size == -1: return (kv_idx < ends2[q_idx]) | (q_idx == kv_idx) else: - return ( - (kv_idx < ends2[q_idx]) - & ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size))) - ) | (q_idx == kv_idx) + return ((kv_idx < ends2[q_idx]) + & ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size)))) | (q_idx == kv_idx) block_mask2 = create_block_mask( attention_mask2, @@ -995,39 +870,23 @@ class CausalMatrixGame2WanModel(BaseDiT): **kwargs, ) -> torch.Tensor: # Extract num_frame_per_block from kwargs - effective_num_frame_per_block = kwargs.pop( - "num_frame_per_block", self.num_frame_per_block - ) + effective_num_frame_per_block = kwargs.pop("num_frame_per_block", self.num_frame_per_block) if mouse_cond is not None or keyboard_cond is not None: - assert ( - self.action_config is not None and len(self.action_config) > 0 - ) + assert (self.action_config is not None and len(self.action_config) > 0) - if encoder_hidden_states is not None and not isinstance( - encoder_hidden_states, torch.Tensor - ): + if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] if isinstance(encoder_hidden_states_image, list): - encoder_hidden_states_image = ( - encoder_hidden_states_image[0] - if len(encoder_hidden_states_image) > 0 - else None - ) + encoder_hidden_states_image = (encoder_hidden_states_image[0] + if len(encoder_hidden_states_image) > 0 else None) ctx = get_forward_context() batch = getattr(ctx, "forward_batch", None) - if ( - batch is not None - and getattr(batch, "image_latent", None) is not None - ): + if (batch is not None and getattr(batch, "image_latent", None) is not None): image_latent = batch.image_latent - if ( - isinstance(image_latent, torch.Tensor) - and image_latent.ndim == 5 - and hidden_states.ndim == 5 - ): + if (isinstance(image_latent, torch.Tensor) and image_latent.ndim == 5 and hidden_states.ndim == 5): _, _, num_frames, _, _ = hidden_states.shape start = start_frame end = start + num_frames @@ -1059,13 +918,9 @@ class CausalMatrixGame2WanModel(BaseDiT): dtype=cond.dtype, ) cond = pad - hidden_states = torch.cat( - [hidden_states, cond.to(dtype=hidden_states.dtype)], dim=1 - ) + hidden_states = torch.cat([hidden_states, cond.to(dtype=hidden_states.dtype)], dim=1) - batch_size, num_channels, num_frames, height, width = ( - hidden_states.shape - ) + batch_size, num_channels, num_frames, height, width = (hidden_states.shape) p_t, p_h, p_w = self.patch_size post_patch_num_frames = num_frames // p_t post_patch_height = height // p_h @@ -1102,20 +957,16 @@ class CausalMatrixGame2WanModel(BaseDiT): timestep_proj, encoder_hidden_states_text, encoder_hidden_states_image, - ) = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image - ) - timestep_proj = timestep_proj.unflatten( - 1, (6, self.hidden_size) - ).unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)) + ) = self.condition_embedder(timestep, encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, + sizes=(batch_size, + post_patch_num_frames)) if encoder_hidden_states_text is not None: if isinstance(encoder_hidden_states_text, list): encoder_hidden_states_text = encoder_hidden_states_text[0] elif encoder_hidden_states_text.ndim == 2: - encoder_hidden_states_text = ( - encoder_hidden_states_text.unsqueeze(0) - ) + encoder_hidden_states_text = (encoder_hidden_states_text.unsqueeze(0)) if encoder_hidden_states_image is not None: if encoder_hidden_states_text is not None: @@ -1172,23 +1023,15 @@ class CausalMatrixGame2WanModel(BaseDiT): for block_index, block in enumerate(self.blocks): causal_kwargs = { "kv_cache": kv_cache[block_index], - "kv_cache_mouse": kv_cache_mouse[block_index] - if kv_cache_mouse - else None, - "kv_cache_keyboard": kv_cache_keyboard[block_index] - if kv_cache_keyboard - else None, - "crossattn_cache": crossattn_cache[block_index] - if crossattn_cache - else None, + "kv_cache_mouse": kv_cache_mouse[block_index] if kv_cache_mouse else None, + "kv_cache_keyboard": kv_cache_keyboard[block_index] if kv_cache_keyboard else None, + "crossattn_cache": crossattn_cache[block_index] if crossattn_cache else None, "current_start": current_start, "cache_start": cache_start, } if torch.is_grad_enabled() and self.gradient_checkpointing: causal_kwargs["num_frame_per_block"] = effective_num_frame_per_block - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, **causal_kwargs - ) + hidden_states = self._gradient_checkpointing_func(block, hidden_states, **causal_kwargs) else: hidden_states = block( hidden_states, @@ -1206,12 +1049,8 @@ class CausalMatrixGame2WanModel(BaseDiT): **causal_kwargs, ) - temb = temb.unflatten( - dim=0, sizes=(batch_size, post_patch_num_frames) - ).unsqueeze(2) - shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk( - 2, dim=2 - ) + temb = temb.unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)).unsqueeze(2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) @@ -1232,33 +1071,19 @@ class CausalMatrixGame2WanModel(BaseDiT): **kwargs, ) -> torch.Tensor: if mouse_cond is not None or keyboard_cond is not None: - assert ( - self.action_config is not None and len(self.action_config) > 0 - ) + assert (self.action_config is not None and len(self.action_config) > 0) - if encoder_hidden_states is not None and not isinstance( - encoder_hidden_states, torch.Tensor - ): + if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] if isinstance(encoder_hidden_states_image, list): - encoder_hidden_states_image = ( - encoder_hidden_states_image[0] - if len(encoder_hidden_states_image) > 0 - else None - ) + encoder_hidden_states_image = (encoder_hidden_states_image[0] + if len(encoder_hidden_states_image) > 0 else None) ctx = get_forward_context() batch = getattr(ctx, "forward_batch", None) - if ( - batch is not None - and getattr(batch, "image_latent", None) is not None - ): + if (batch is not None and getattr(batch, "image_latent", None) is not None): image_latent = batch.image_latent - if ( - isinstance(image_latent, torch.Tensor) - and image_latent.ndim == 5 - and hidden_states.ndim == 5 - ): + if (isinstance(image_latent, torch.Tensor) and image_latent.ndim == 5 and hidden_states.ndim == 5): _, _, num_frames, _, _ = hidden_states.shape cond = image_latent if cond.shape[2] >= num_frames: @@ -1275,13 +1100,9 @@ class CausalMatrixGame2WanModel(BaseDiT): dtype=cond.dtype, ) cond = torch.cat([cond, pad], dim=2) - hidden_states = torch.cat( - [hidden_states, cond.to(dtype=hidden_states.dtype)], dim=1 - ) + hidden_states = torch.cat([hidden_states, cond.to(dtype=hidden_states.dtype)], dim=1) - batch_size, num_channels, num_frames, height, width = ( - hidden_states.shape - ) + batch_size, num_channels, num_frames, height, width = (hidden_states.shape) p_t, p_h, p_w = self.patch_size post_patch_num_frames: int = num_frames // p_t post_patch_height: int = height // p_h @@ -1355,20 +1176,16 @@ class CausalMatrixGame2WanModel(BaseDiT): timestep_proj, encoder_hidden_states_text, encoder_hidden_states_image, - ) = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image - ) - timestep_proj = timestep_proj.unflatten( - 1, (6, self.hidden_size) - ).unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)) + ) = self.condition_embedder(timestep, encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, + sizes=(batch_size, + post_patch_num_frames)) if encoder_hidden_states_text is not None: if isinstance(encoder_hidden_states_text, list): encoder_hidden_states_text = encoder_hidden_states_text[0] elif encoder_hidden_states_text.ndim == 2: - encoder_hidden_states_text = ( - encoder_hidden_states_text.unsqueeze(0) - ) + encoder_hidden_states_text = (encoder_hidden_states_text.unsqueeze(0)) if encoder_hidden_states_image is not None: if encoder_hidden_states_text is not None: @@ -1398,12 +1215,8 @@ class CausalMatrixGame2WanModel(BaseDiT): # Initialize kv_cache, kv_cache_mouse, kv_cache_keyboard, crossattn_cache if not provided kv_cache = kwargs.pop("kv_cache", [None] * len(self.blocks)) kv_cache_mouse = kwargs.pop("kv_cache_mouse", [None] * len(self.blocks)) - kv_cache_keyboard = kwargs.pop( - "kv_cache_keyboard", [None] * len(self.blocks) - ) - crossattn_cache = kwargs.pop( - "crossattn_cache", [None] * len(self.blocks) - ) + kv_cache_keyboard = kwargs.pop("kv_cache_keyboard", [None] * len(self.blocks)) + crossattn_cache = kwargs.pop("crossattn_cache", [None] * len(self.blocks)) for i, block in enumerate(self.blocks): block_crossattn_cache = None @@ -1452,12 +1265,8 @@ class CausalMatrixGame2WanModel(BaseDiT): use_rope_keyboard=self.use_rope_keyboard, ) - temb = temb.unflatten( - dim=0, sizes=(batch_size, post_patch_num_frames) - ).unsqueeze(2) - shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk( - 2, dim=2 - ) + temb = temb.unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)).unsqueeze(2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) @@ -1476,7 +1285,7 @@ class CausalMatrixGame2WanModel(BaseDiT): p_t, p_h, p_w = self.patch_size f, h, w = grid_sizes - x = x[:, : f * h * w].view(-1, f, h, w, p_t, p_h, p_w, c) + x = x[:, :f * h * w].view(-1, f, h, w, p_t, p_h, p_w, c) x = x.permute(0, 7, 1, 4, 2, 5, 3, 6) x = x.reshape(-1, c, f * p_t, h * p_h, w * p_w) return x diff --git a/fastvideo/models/dits/matrixgame2/model.py b/fastvideo/models/dits/matrixgame2/model.py index 2bf90b93b..f3c680a9d 100644 --- a/fastvideo/models/dits/matrixgame2/model.py +++ b/fastvideo/models/dits/matrixgame2/model.py @@ -44,6 +44,7 @@ logger = init_logger(__name__) class MatrixGame2TimeImageEmbedding(nn.Module): + def __init__( self, dim: int, @@ -52,12 +53,8 @@ class MatrixGame2TimeImageEmbedding(nn.Module): ): super().__init__() - self.time_embedder = TimestepEmbedder( - dim, frequency_embedding_size=time_freq_dim, act_layer="silu" - ) - self.time_modulation = ModulateProjection( - dim, factor=6, act_layer="silu" - ) + self.time_embedder = TimestepEmbedder(dim, frequency_embedding_size=time_freq_dim, act_layer="silu") + self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") self.image_embedder = None if image_embed_dim is not None: @@ -77,9 +74,7 @@ class MatrixGame2TimeImageEmbedding(nn.Module): if encoder_hidden_states_image is not None: assert self.image_embedder is not None - encoder_hidden_states_image = self.image_embedder( - encoder_hidden_states_image - ) + encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image) encoder_hidden_states = torch.zeros( (timestep.shape[0], 0, temb.shape[-1]), @@ -96,6 +91,7 @@ class MatrixGame2TimeImageEmbedding(nn.Module): class MatrixGame2CrossAttention(WanSelfAttention): + def forward(self, x, context, context_lens=None, crossattn_cache=None): r""" Args: @@ -133,6 +129,7 @@ class MatrixGame2CrossAttention(WanSelfAttention): class MatrixGame2TransformerBlock(nn.Module): + def __init__( self, dim: int, @@ -190,14 +187,10 @@ class MatrixGame2TransformerBlock(nn.Module): # 2. Cross-attention if added_kv_proj_dim is not None: # I2V - self.attn2 = WanI2VCrossAttention( - dim, num_heads, qk_norm=qk_norm, eps=eps - ) + self.attn2 = WanI2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) else: # T2V - self.attn2 = WanT2VCrossAttention( - dim, num_heads, qk_norm=qk_norm, eps=eps - ) + self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, norm_type="layer", @@ -244,9 +237,7 @@ class MatrixGame2TransformerBlock(nn.Module): c_shift_msa, c_scale_msa, c_gate_msa, - ) = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk( - 6, dim=2 - ) + ) = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2) # batch_size, seq_len, 1, inner_dim shift_msa = shift_msa.squeeze(2) scale_msa = scale_msa.squeeze(2) @@ -268,9 +259,7 @@ class MatrixGame2TransformerBlock(nn.Module): assert shift_msa.dtype == torch.float32 # 1. Self-attention - norm_hidden_states = ( - self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa - ).to(orig_dtype) + norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -297,21 +286,17 @@ class MatrixGame2TransformerBlock(nn.Module): attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.tensor([0], device=hidden_states.device) - norm_hidden_states, hidden_states = self.self_attn_residual_norm( - hidden_states, attn_output, gate_msa, null_shift, null_scale - ) + norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, + null_shift, null_scale) norm_hidden_states, hidden_states = ( norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype), ) # 2. Cross-attention - attn_output = self.attn2( - norm_hidden_states, context=encoder_hidden_states, context_lens=None - ) - norm_hidden_states, hidden_states = self.cross_attn_residual_norm( - hidden_states, attn_output, 1, c_shift_msa, c_scale_msa - ) + attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None) + norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa, + c_scale_msa) norm_hidden_states, hidden_states = ( norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype), @@ -350,16 +335,10 @@ class MatrixGame2WanModel(BaseDiT): _fsdp_shard_conditions = _DEFAULT_MATRIXGAME2_CONFIG._fsdp_shard_conditions _compile_conditions = _DEFAULT_MATRIXGAME2_CONFIG._compile_conditions - _supported_attention_backends = ( - _DEFAULT_MATRIXGAME2_CONFIG._supported_attention_backends - ) + _supported_attention_backends = (_DEFAULT_MATRIXGAME2_CONFIG._supported_attention_backends) param_names_mapping = _DEFAULT_MATRIXGAME2_CONFIG.param_names_mapping - reverse_param_names_mapping = ( - _DEFAULT_MATRIXGAME2_CONFIG.reverse_param_names_mapping - ) - lora_param_names_mapping = ( - _DEFAULT_MATRIXGAME2_CONFIG.lora_param_names_mapping - ) + reverse_param_names_mapping = (_DEFAULT_MATRIXGAME2_CONFIG.reverse_param_names_mapping) + lora_param_names_mapping = (_DEFAULT_MATRIXGAME2_CONFIG.lora_param_names_mapping) def __init__( self, @@ -394,23 +373,20 @@ class MatrixGame2WanModel(BaseDiT): self.action_config = getattr(config, "action_config", {}) # 3. Transformer blocks - self.blocks = nn.ModuleList( - [ - MatrixGame2TransformerBlock( - inner_dim, - config.ffn_dim, - config.num_attention_heads, - config.qk_norm, - config.cross_attn_norm, - config.eps, - config.added_kv_proj_dim, - self._supported_attention_backends, - prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", - action_config=self.action_config, - ) - for i in range(config.num_layers) - ] - ) + self.blocks = nn.ModuleList([ + MatrixGame2TransformerBlock( + inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + self._supported_attention_backends, + prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", + action_config=self.action_config, + ) for i in range(config.num_layers) + ]) # 4. Output norm & projection self.norm_out = LayerNormScaleShift( @@ -421,12 +397,8 @@ class MatrixGame2WanModel(BaseDiT): dtype=torch.float32, compute_dtype=torch.float32, ) - self.proj_out = nn.Linear( - inner_dim, config.out_channels * math.prod(config.patch_size) - ) - self.scale_shift_table = nn.Parameter( - torch.randn(1, 2, inner_dim) / inner_dim**0.5 - ) + self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) self.gradient_checkpointing = False @@ -443,21 +415,14 @@ class MatrixGame2WanModel(BaseDiT): keyboard_cond: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: - if encoder_hidden_states is not None and not isinstance( - encoder_hidden_states, torch.Tensor - ): + if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] - if ( - isinstance(encoder_hidden_states_image, list) - and len(encoder_hidden_states_image) > 0 - ): + if (isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0): encoder_hidden_states_image = encoder_hidden_states_image[0] elif not isinstance(encoder_hidden_states_image, torch.Tensor): encoder_hidden_states_image = None - batch_size, num_channels, num_frames, height, width = ( - hidden_states.shape - ) + batch_size, num_channels, num_frames, height, width = (hidden_states.shape) p_t, p_h, p_w = self.patch_size post_patch_num_frames = num_frames // p_t post_patch_height = height // p_h @@ -481,11 +446,7 @@ class MatrixGame2WanModel(BaseDiT): ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = ( - (freqs_cos.float(), freqs_sin.float()) - if freqs_cos is not None - else None - ) + freqs_cis = ((freqs_cos.float(), freqs_sin.float()) if freqs_cos is not None else None) hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) @@ -498,9 +459,7 @@ class MatrixGame2WanModel(BaseDiT): timestep_proj, encoder_hidden_states, encoder_hidden_states_image, - ) = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image - ) + ) = self.condition_embedder(timestep, encoder_hidden_states, encoder_hidden_states_image) timestep_proj = timestep_proj.unflatten(1, (6, -1)) if encoder_hidden_states is not None: @@ -515,9 +474,7 @@ class MatrixGame2WanModel(BaseDiT): if encoder_hidden_states_image is not None: if encoder_hidden_states is not None: - encoder_hidden_states = torch.concat( - [encoder_hidden_states_image, encoder_hidden_states], dim=1 - ) + encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) else: encoder_hidden_states = encoder_hidden_states_image @@ -552,9 +509,7 @@ class MatrixGame2WanModel(BaseDiT): ) # Output - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk( - 2, dim=1 - ) + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) diff --git a/fastvideo/models/dits/matrixgame2/utils.py b/fastvideo/models/dits/matrixgame2/utils.py index f51f7251e..71f6981a2 100644 --- a/fastvideo/models/dits/matrixgame2/utils.py +++ b/fastvideo/models/dits/matrixgame2/utils.py @@ -16,19 +16,20 @@ try: except ImportError: cv2 = None - CAM_VALUE = 0.1 -CAMERA_MAP = { - "i": [CAM_VALUE, 0], "k": [-CAM_VALUE, 0], - "j": [0, -CAM_VALUE], "l": [0, CAM_VALUE], "u": [0, 0] -} +CAMERA_MAP = {"i": [CAM_VALUE, 0], "k": [-CAM_VALUE, 0], "j": [0, -CAM_VALUE], "l": [0, CAM_VALUE], "u": [0, 0]} KEYBOARD_MAP_4 = { # base_distilled_model (universal): W/S/A/D - "w": [1, 0, 0, 0], "s": [0, 1, 0, 0], - "a": [0, 0, 1, 0], "d": [0, 0, 0, 1], "q": [0, 0, 0, 0] + "w": [1, 0, 0, 0], + "s": [0, 1, 0, 0], + "a": [0, 0, 1, 0], + "d": [0, 0, 0, 1], + "q": [0, 0, 0, 0] } KEYBOARD_MAP_2 = { # gta_distilled_model: W/S only (steering via mouse) - "w": [1, 0], "s": [0, 1], "q": [0, 0] + "w": [1, 0], + "s": [0, 1], + "q": [0, 0] } KEYBOARD_MAP_7 = { # templerun_distilled_model: still/w/s/left/right/a/d "q": [1, 0, 0, 0, 0, 0, 0], # still @@ -71,24 +72,17 @@ def retain_kv_with_sink( def _require_cv2() -> bool: if cv2 is None: - logger.warning( - "OpenCV is not available; skipping MatrixGame2 validation overlay." - ) + logger.warning("OpenCV is not available; skipping MatrixGame2 validation overlay.") return False return True -def _matrixgame2_overlay_keys_from_keyboard( - keyboard_frame: np.ndarray | torch.Tensor, -) -> dict[str, bool]: +def _matrixgame2_overlay_keys_from_keyboard(keyboard_frame: np.ndarray | torch.Tensor, ) -> dict[str, bool]: vec = np.asarray(keyboard_frame, dtype=np.float32).reshape(-1) dim = vec.shape[0] if dim >= 23: - return { - key: bool(vec[idx] > 0.5) - for key, idx in SOLARIS_MOVEMENT_KEY_INDICES.items() - } + return {key: bool(vec[idx] > 0.5) for key, idx in SOLARIS_MOVEMENT_KEY_INDICES.items()} if dim == 7: return { "W": bool(vec[1] > 0.5), @@ -145,10 +139,10 @@ def draw_rounded_rectangle( def draw_keys_on_frame( - frame: np.ndarray, - keys: dict[str, bool], - key_size: tuple[int, int] = (30, 30), - top_margin: int = 15, + frame: np.ndarray, + keys: dict[str, bool], + key_size: tuple[int, int] = (30, 30), + top_margin: int = 15, ) -> None: if not _require_cv2(): return @@ -177,9 +171,7 @@ def draw_keys_on_frame( radius=5, alpha=alpha, ) - text_size = cv2.getTextSize( - key, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1 - )[0] + text_size = cv2.getTextSize(key, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0] text_x = x + (key_size[0] - text_size[0]) // 2 text_y = y + (key_size[1] + text_size[1]) // 2 cv2.putText( @@ -295,38 +287,37 @@ def expand_action_to_frames(action: dict, num_frames: int) -> tuple[torch.Tensor result = {} for key, tensor in action.items(): if tensor is not None: - # Expand to [num_frames, D] then unsqueeze to [1, num_frames, D] + # Expand to [num_frames, D] then unsqueeze to [1, num_frames, D] result[key] = tensor.unsqueeze(0).repeat(num_frames, 1).unsqueeze(0) else: result[key] = None if "mouse" not in result or result["mouse"] is None: - # keyboard device if available, otherwise default - device = result.get("keyboard", torch.tensor([])).device if result.get("keyboard") is not None else get_local_torch_device() - result["mouse"] = torch.zeros(1, num_frames, 2, device=device) - + # keyboard device if available, otherwise default + device = result.get("keyboard", torch.tensor( + [])).device if result.get("keyboard") is not None else get_local_torch_device() + result["mouse"] = torch.zeros(1, num_frames, 2, device=device) + return result["keyboard"], result["mouse"] + def get_current_action(mode="universal"): CAM_VALUE = 0.1 if mode == 'universal': logger.info("") - logger.info('-'*30) + logger.info('-' * 30) logger.info("PRESS [I, K, J, L, U] FOR CAMERA TRANSFORM\n (I: up, K: down, J: left, L: right, U: no move)") logger.info("PRESS [W, S, A, D, Q] FOR MOVEMENT\n (W: forward, S: back, A: left, D: right, Q: no move)") - logger.info('-'*30) + logger.info('-' * 30) CAMERA_VALUE_MAP = { - "i": [CAM_VALUE, 0], - "k": [-CAM_VALUE, 0], - "j": [0, -CAM_VALUE], - "l": [0, CAM_VALUE], - "u": [0, 0] - } - KEYBOARD_IDX = { - "w": [1, 0, 0, 0], "s": [0, 1, 0, 0], "a": [0, 0, 1, 0], "d": [0, 0, 0, 1], - "q": [0, 0, 0, 0] + "i": [CAM_VALUE, 0], + "k": [-CAM_VALUE, 0], + "j": [0, -CAM_VALUE], + "l": [0, CAM_VALUE], + "u": [0, 0] } + KEYBOARD_IDX = {"w": [1, 0, 0, 0], "s": [0, 1, 0, 0], "a": [0, 0, 1, 0], "d": [0, 0, 0, 1], "q": [0, 0, 0, 0]} flag = 0 while flag != 1: try: @@ -340,22 +331,17 @@ def get_current_action(mode="universal"): keyboard_cond = torch.tensor(KEYBOARD_IDX[idx_keyboard]).cuda() elif mode == 'gta_drive': logger.info("") - logger.info('-'*30) + logger.info('-' * 30) logger.info("PRESS [W, S, A, D, Q] FOR MOVEMENT\n (W: forward, S: back, A: left, D: right, Q: no move)") - logger.info('-'*30) - CAMERA_VALUE_MAP = { - "a": [0, -CAM_VALUE], - "d": [0, CAM_VALUE], - "q": [0, 0] - } - KEYBOARD_IDX = { - "w": [1, 0], "s": [0, 1], - "q": [0, 0] - } + logger.info('-' * 30) + CAMERA_VALUE_MAP = {"a": [0, -CAM_VALUE], "d": [0, CAM_VALUE], "q": [0, 0]} + KEYBOARD_IDX = {"w": [1, 0], "s": [0, 1], "q": [0, 0]} flag = 0 while flag != 1: try: - indexes = input('Please input the actions (split with ` `):\n(e.g. `W` for forward, `W A` for forward and left)\n').strip().lower().split(' ') + indexes = input( + 'Please input the actions (split with ` `):\n(e.g. `W` for forward, `W A` for forward and left)\n' + ).strip().lower().split(' ') idx_mouse = [] idx_keyboard = [] for i in indexes: @@ -375,43 +361,47 @@ def get_current_action(mode="universal"): keyboard_cond = torch.tensor(KEYBOARD_IDX[idx_keyboard[0]]).cuda() elif mode == 'templerun': logger.info("") - logger.info('-'*30) - logger.info("PRESS [W, S, A, D, Z, C, Q] FOR ACTIONS\n (W: jump, S: slide, A: left side, D: right side, Z: turn left, C: turn right, Q: no move)") - logger.info('-'*30) - KEYBOARD_IDX = { - "w": [0, 1, 0, 0, 0, 0, 0], "s": [0, 0, 1, 0, 0, 0, 0], - "a": [0, 0, 0, 0, 0, 1, 0], "d": [0, 0, 0, 0, 0, 0, 1], - "z": [0, 0, 0, 1, 0, 0, 0], "c": [0, 0, 0, 0, 1, 0, 0], + logger.info('-' * 30) + logger.info( + "PRESS [W, S, A, D, Z, C, Q] FOR ACTIONS\n (W: jump, S: slide, A: left side, D: right side, Z: turn left, C: turn right, Q: no move)" + ) + logger.info('-' * 30) + KEYBOARD_IDX = { + "w": [0, 1, 0, 0, 0, 0, 0], + "s": [0, 0, 1, 0, 0, 0, 0], + "a": [0, 0, 0, 0, 0, 1, 0], + "d": [0, 0, 0, 0, 0, 0, 1], + "z": [0, 0, 0, 1, 0, 0, 0], + "c": [0, 0, 0, 0, 1, 0, 0], "q": [1, 0, 0, 0, 0, 0, 0] } flag = 0 while flag != 1: try: - idx_keyboard = input('Please input the action: \n(e.g. `W` for forward, `Z` for turning left)\n').strip().lower() + idx_keyboard = input( + 'Please input the action: \n(e.g. `W` for forward, `Z` for turning left)\n').strip().lower() if idx_keyboard in KEYBOARD_IDX.keys(): flag = 1 except Exception: pass keyboard_cond = torch.tensor(KEYBOARD_IDX[idx_keyboard]).cuda() - + if mode != 'templerun': - return { - "mouse": mouse_cond, - "keyboard": keyboard_cond - } - return { - "keyboard": keyboard_cond - } + return {"mouse": mouse_cond, "keyboard": keyboard_cond} + return {"keyboard": keyboard_cond} + async def get_current_action_async(mode="universal"): return await asyncio.to_thread(get_current_action, mode) + def load_initial_image(image_path: str = None) -> Image.Image: if image_path and os.path.exists(image_path): return Image.open(image_path).convert("RGB") logger.warning("No image provided, creating placeholder...") return Image.new("RGB", (640, 352), (128, 128, 128)) + def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = None): if keyboard_dim not in (2, 4, 6, 7): raise ValueError(f"keyboard_dim must be 2, 4, 6, or 7, got {keyboard_dim}") @@ -423,7 +413,7 @@ def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = No random.seed(seed) num_samples_per_action = 4 - + # Define actions based on keyboard_dim if keyboard_dim == 4: # Universal model: W, S, A, D @@ -449,9 +439,7 @@ def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = No actions_single_camera = [] # No mouse for Temple Run keyboard_idx = {"still": 0, "forward": 1, "back": 2, "left": 3, "right": 4, "a": 5, "d": 6} - actions_to_test = ( - actions_double_action * 5 + actions_single_camera * 5 + actions_single_action * 5 - ) + actions_to_test = (actions_double_action * 5 + actions_single_camera * 5 + actions_single_action * 5) for action in (actions_single_action + actions_double_action): for camera in actions_single_camera: actions_to_test.append(f"{action}_{camera}") @@ -518,6 +506,7 @@ def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = No return {"keyboard": keyboard_condition, "mouse": mouse_condition} + def parse_config(config, mode="universal"): assert mode in ['universal', 'gta_drive', 'templerun'] key_data = {} @@ -555,6 +544,7 @@ def parse_config(config, mode="universal"): ) return key_data, mouse_data + # NOTE: drawing functions are commented out to avoid cv2/libGL dependency. # # def draw_rounded_rectangle(image, top_left, bottom_right, color, radius=10, alpha=0.5): @@ -576,21 +566,21 @@ def parse_config(config, mode="universal"): # vertical_shift = -20 # horizon_shift_all = 50 # key_positions = { -# "W": (w // 2 - key_size[0] // 2 - horison_shift - horizon_shift_all, +# "W": (w // 2 - key_size[0] // 2 - horison_shift - horizon_shift_all, # h - bottom_margin - key_size[1] * 2 + vertical_shift - 20), -# "A": (w // 2 - key_size[0] * 2 + 5 - horison_shift - horizon_shift_all, +# "A": (w // 2 - key_size[0] * 2 + 5 - horison_shift - horizon_shift_all, # h - bottom_margin - key_size[1] + vertical_shift), -# "S": (w // 2 - key_size[0] // 2 - horison_shift - horizon_shift_all, +# "S": (w // 2 - key_size[0] // 2 - horison_shift - horizon_shift_all, # h - bottom_margin - key_size[1] + vertical_shift), -# "D": (w // 2 + key_size[0] - 5 - horison_shift - horizon_shift_all, +# "D": (w // 2 + key_size[0] - 5 - horison_shift - horizon_shift_all, # h - bottom_margin - key_size[1] + vertical_shift), # } # key_icon = {"W": "W", "A": "A", "S": "S", "D": "D", "left": "left", "right": "right"} # if mode == 'templerun': # key_positions.update({ -# "left": (w // 2 + key_size[0] * 2 + spacing * 2 - horison_shift - horizon_shift_all, +# "left": (w // 2 + key_size[0] * 2 + spacing * 2 - horison_shift - horizon_shift_all, # h - bottom_margin - key_size[1] + vertical_shift), -# "right": (w // 2 + key_size[0] * 3 + spacing * 7 - horison_shift - horizon_shift_all, +# "right": (w // 2 + key_size[0] * 3 + spacing * 7 - horison_shift - horizon_shift_all, # h - bottom_margin - key_size[1] + vertical_shift) # }) # @@ -651,7 +641,7 @@ def parse_config(config, mode="universal"): # frame_region[:, :, c] = (1 - alpha) * frame_region[:, :, c] + alpha * icon_rgb[:, :, c] # frame[top_left_y:bottom_right_y, top_left_x:bottom_right_x] = frame_region # -# def process_video(input_video, output_video, config, mouse_icon_path, +# def process_video(input_video, output_video, config, mouse_icon_path, # mouse_scale=1.0, mouse_rotation=0, process_icon=True, mode='universal'): # key_data, mouse_data = parse_config(config, mode=mode) # fps = 12 @@ -670,6 +660,6 @@ def parse_config(config, mode="universal"): # mouse_position = mouse_data.get(frame_idx, (frame_width // 2, frame_height // 2)) # overlay_icon(frame, mouse_icon, mouse_position, scale=mouse_scale, rotation=mouse_rotation) # out_video.append(frame / 255) -# +# # export_to_video(out_video, output_video, fps=fps) # logger.info(f"Video saved to {output_video}") diff --git a/fastvideo/models/dits/matrixgame3/action_module.py b/fastvideo/models/dits/matrixgame3/action_module.py index 126ac815a..bd1fa4aec 100644 --- a/fastvideo/models/dits/matrixgame3/action_module.py +++ b/fastvideo/models/dits/matrixgame3/action_module.py @@ -90,12 +90,8 @@ class MatrixGame3ActionModule(nn.Module): ): super().__init__() patch_size = patch_size if patch_size is not None else [1, 2, 2] - rope_dim_list = ( - rope_dim_list if rope_dim_list is not None else [8, 28, 28] - ) - mouse_qk_dim_list = ( - mouse_qk_dim_list if mouse_qk_dim_list is not None else [8, 28, 28] - ) + rope_dim_list = (rope_dim_list if rope_dim_list is not None else [8, 28, 28]) + mouse_qk_dim_list = (mouse_qk_dim_list if mouse_qk_dim_list is not None else [8, 28, 28]) blocks = blocks if blocks is not None else [] self.enable_mouse = enable_mouse self.enable_keyboard = enable_keyboard @@ -115,8 +111,7 @@ class MatrixGame3ActionModule(nn.Module): c = mouse_hidden_dim self.mouse_mlp = nn.Sequential( nn.Linear( - mouse_dim_in * vae_time_compression_ratio * windows_size - + img_hidden_size, + mouse_dim_in * vae_time_compression_ratio * windows_size + img_hidden_size, c, bias=True, ), @@ -127,64 +122,42 @@ class MatrixGame3ActionModule(nn.Module): head_dim = c // heads_num self.t_qkv = ReplicatedLinear(c, c * 3, bias=qkv_bias) - self.img_attn_q_norm = ( - WanRMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity() - ) - self.img_attn_k_norm = ( - WanRMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity() - ) - self.proj_mouse = ReplicatedLinear( - c, img_hidden_size, bias=qkv_bias - ) + self.img_attn_q_norm = (WanRMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity()) + self.img_attn_k_norm = (WanRMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity()) + self.proj_mouse = ReplicatedLinear(c, img_hidden_size, bias=qkv_bias) if self.enable_keyboard: head_dim_key = keyboard_hidden_dim // heads_num - self.key_attn_q_norm = ( - WanRMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity() - ) - self.key_attn_k_norm = ( - WanRMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity() - ) + self.key_attn_q_norm = (WanRMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity()) + self.key_attn_k_norm = (WanRMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity()) - self.mouse_attn_q = ReplicatedLinear( - img_hidden_size, keyboard_hidden_dim, bias=qkv_bias - ) + self.mouse_attn_q = ReplicatedLinear(img_hidden_size, keyboard_hidden_dim, bias=qkv_bias) self.keyboard_attn_kv = ReplicatedLinear( hidden_size * windows_size * vae_time_compression_ratio, keyboard_hidden_dim * 2, bias=qkv_bias, ) - self.proj_keyboard = ReplicatedLinear( - keyboard_hidden_dim, img_hidden_size, bias=qkv_bias - ) + self.proj_keyboard = ReplicatedLinear(keyboard_hidden_dim, img_hidden_size, bias=qkv_bias) - self.mouse_attn_layer = ( - LocalAttention( - num_heads=heads_num, - head_size=mouse_hidden_dim // heads_num, - causal=False, - supported_attention_backends=( - AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA, - ), - ) - if self.enable_mouse - else None - ) + self.mouse_attn_layer = (LocalAttention( + num_heads=heads_num, + head_size=mouse_hidden_dim // heads_num, + causal=False, + supported_attention_backends=( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ), + ) if self.enable_mouse else None) - self.keyboard_attn_layer = ( - LocalAttention( - num_heads=heads_num, - head_size=keyboard_hidden_dim // heads_num, - causal=False, - supported_attention_backends=( - AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA, - ), - ) - if self.enable_keyboard - else None - ) + self.keyboard_attn_layer = (LocalAttention( + num_heads=heads_num, + head_size=keyboard_hidden_dim // heads_num, + causal=False, + supported_attention_backends=( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ), + ) if self.enable_keyboard else None) self.vae_time_compression_ratio = vae_time_compression_ratio self.windows_size = windows_size @@ -236,33 +209,20 @@ class MatrixGame3ActionModule(nn.Module): if isinstance(self.patch_size, int): assert all(s % self.patch_size == 0 for s in latents_size), ( f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), " - f"but got {latents_size}." - ) + f"but got {latents_size}.") rope_sizes = [s // self.patch_size for s in latents_size] elif isinstance(self.patch_size, list): - assert all( - s % self.patch_size[idx] == 0 - for idx, s in enumerate(latents_size) - ), ( + assert all(s % self.patch_size[idx] == 0 for idx, s in enumerate(latents_size)), ( f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), " - f"but got {latents_size}." - ) - rope_sizes = [ - s // self.patch_size[idx] for idx, s in enumerate(latents_size) - ] + f"but got {latents_size}.") + rope_sizes = [s // self.patch_size[idx] for idx, s in enumerate(latents_size)] if len(rope_sizes) != target_ndim: - rope_sizes = [1] * ( - target_ndim - len(rope_sizes) - ) + rope_sizes # time axis + rope_sizes = [1] * (target_ndim - len(rope_sizes)) + rope_sizes # time axis if rope_dim_list is None: - rope_dim_list = [ - head_dim // target_ndim for _ in range(target_ndim) - ] - assert sum(rope_dim_list) == head_dim, ( - "sum(rope_dim_list) should equal to head_dim of attention layer" - ) + rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] + assert sum(rope_dim_list) == head_dim, ("sum(rope_dim_list) should equal to head_dim of attention layer") # Use Matrix-Game wrapper for FastVideo's function freqs_cos, freqs_sin = _get_nd_rotary_pos_embed_matrixgame( rope_dim_list, @@ -270,17 +230,9 @@ class MatrixGame3ActionModule(nn.Module): theta=self.rope_theta, theta_rescale_factor=1, ) - return freqs_cos[ - -video_length - * rope_sizes[1] - * rope_sizes[2] - // self.patch_size[0] : - ], freqs_sin[ - -video_length - * rope_sizes[1] - * rope_sizes[2] - // self.patch_size[0] : - ] + return freqs_cos[-video_length * rope_sizes[1] * rope_sizes[2] // + self.patch_size[0]:], freqs_sin[-video_length * rope_sizes[1] * rope_sizes[2] // + self.patch_size[0]:] def _forward_mouse( self, @@ -301,11 +253,10 @@ class MatrixGame3ActionModule(nn.Module): group_mouse = [ mouse_condition[ :, - self.vae_time_compression_ratio * (i - self.windows_size) - + pad_t : i * self.vae_time_compression_ratio + pad_t, + self.vae_time_compression_ratio * (i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + + pad_t, :, - ] - for i in range(N_feats) + ] for i in range(N_feats) ] group_mouse = torch.stack(group_mouse, dim=1) mem_len = 0 @@ -317,19 +268,13 @@ class MatrixGame3ActionModule(nn.Module): actual_num_frames = group_mouse.shape[1] S = th * tw - group_mouse = group_mouse.unsqueeze(-1).expand( - B, actual_num_frames, pad_t, C, S - ) - group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape( - B * S, actual_num_frames, pad_t * C - ) + group_mouse = group_mouse.unsqueeze(-1).expand(B, actual_num_frames, pad_t, C, S) + group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape(B * S, actual_num_frames, pad_t * C) group_mouse = torch.cat([hidden_states, group_mouse], dim=-1) group_mouse = self.mouse_mlp(group_mouse) mouse_qkv, _ = self.t_qkv(group_mouse) - q, k, v = rearrange( - mouse_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num - ) + q, k, v = rearrange(mouse_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) q = self.img_attn_q_norm(q).to(v) k = self.img_attn_k_norm(k).to(v) if mem_len > 0: @@ -364,25 +309,20 @@ class MatrixGame3ActionModule(nn.Module): group_keyboard = [ keyboard_condition[ :, - self.vae_time_compression_ratio * (i - self.windows_size) - + pad_t : i * self.vae_time_compression_ratio + pad_t, + self.vae_time_compression_ratio * (i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + + pad_t, :, - ] - for i in range(N_feats) + ] for i in range(N_feats) ] group_keyboard = torch.stack(group_keyboard, dim=1) mem_len = 0 if keyboard_cond_memory is not None: mem_len = keyboard_cond_memory.shape[1] - k_mem = keyboard_cond_memory.to( - device=group_keyboard.device, dtype=group_keyboard.dtype - ) + k_mem = keyboard_cond_memory.to(device=group_keyboard.device, dtype=group_keyboard.dtype) k_mem = self.keyboard_embed(k_mem) k_mem = k_mem.unsqueeze(2).expand(-1, -1, pad_t, -1) group_keyboard = torch.cat([k_mem, group_keyboard], dim=1) - group_keyboard = group_keyboard.reshape( - shape=(group_keyboard.shape[0], group_keyboard.shape[1], -1) - ) + group_keyboard = group_keyboard.reshape(shape=(group_keyboard.shape[0], group_keyboard.shape[1], -1)) mouse_q, _ = self.mouse_attn_q(hidden_states) keyboard_kv, _ = self.keyboard_attn_kv(group_keyboard) @@ -392,9 +332,7 @@ class MatrixGame3ActionModule(nn.Module): q = mouse_q.view(B, L, self.heads_num, D) B, L, KHD = keyboard_kv.shape - k, v = keyboard_kv.view(B, L, 2, self.heads_num, D).permute( - 2, 0, 1, 3, 4 - ) + k, v = keyboard_kv.view(B, L, 2, self.heads_num, D).permute(2, 0, 1, 3, 4) q = self.key_attn_q_norm(q).to(v) k = self.key_attn_k_norm(k).to(v) @@ -438,30 +376,23 @@ class MatrixGame3ActionModule(nn.Module): target_device = x.device target_dtype = x.dtype if mouse_condition is not None: - mouse_condition = mouse_condition.to( - device=target_device, dtype=target_dtype - ) + mouse_condition = mouse_condition.to(device=target_device, dtype=target_dtype) if keyboard_condition is not None: - keyboard_condition = keyboard_condition.to( - device=target_device, dtype=target_dtype - ) + keyboard_condition = keyboard_condition.to(device=target_device, dtype=target_dtype) else: return x B, N_frames, C = keyboard_condition.shape assert tt * th * tw == x.shape[1] - assert ( - ((N_frames - 1) + self.vae_time_compression_ratio) % self.vae_time_compression_ratio == 0 - or (N_frames % self.vae_time_compression_ratio == 0) - ) + assert (((N_frames - 1) + self.vae_time_compression_ratio) % self.vae_time_compression_ratio == 0 + or (N_frames % self.vae_time_compression_ratio == 0)) if ((N_frames - 1) + self.vae_time_compression_ratio) % self.vae_time_compression_ratio == 0: N_feats = int((N_frames - 1) / self.vae_time_compression_ratio) + 1 else: N_feats = N_frames // self.vae_time_compression_ratio # Lazy initialization of freqs - if (self._freqs_cos is None or self._freqs_sin is None - or self._freqs_cos.device != x.device): + if (self._freqs_cos is None or self._freqs_sin is None or self._freqs_cos.device != x.device): fc, fs = self.get_rotary_pos_embed( 7500, self.patch_size[1], @@ -475,9 +406,7 @@ class MatrixGame3ActionModule(nn.Module): freqs_cis = (self._freqs_cos, self._freqs_sin) if self.enable_mouse and mouse_condition is not None: - hidden_states = rearrange( - x, "B (T S) C -> (B S) T C", T=tt, S=th * tw - ) + hidden_states = rearrange(x, "B (T S) C -> (B S) T C", T=tt, S=th * tw) B, N_frames, C = mouse_condition.shape else: hidden_states = x diff --git a/fastvideo/models/dits/matrixgame3/model.py b/fastvideo/models/dits/matrixgame3/model.py index bd800c558..2d1e64749 100644 --- a/fastvideo/models/dits/matrixgame3/model.py +++ b/fastvideo/models/dits/matrixgame3/model.py @@ -138,6 +138,7 @@ def _apply_rope_with_frame_indices( class MatrixGame3TimeImageEmbedding(nn.Module): + def __init__( self, dim: int, @@ -145,12 +146,8 @@ class MatrixGame3TimeImageEmbedding(nn.Module): ): super().__init__() - self.time_embedder = TimestepEmbedder( - dim, frequency_embedding_size=time_freq_dim, act_layer="silu" - ) - self.time_modulation = ModulateProjection( - dim, factor=6, act_layer="silu" - ) + self.time_embedder = TimestepEmbedder(dim, frequency_embedding_size=time_freq_dim, act_layer="silu") + self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") def forward( self, @@ -173,6 +170,7 @@ class MatrixGame3TimeImageEmbedding(nn.Module): class MatrixGame3CrossAttention(WanSelfAttention): + def forward(self, x, context, context_lens=None): r""" Args: @@ -197,6 +195,7 @@ class MatrixGame3CrossAttention(WanSelfAttention): class MatrixGame3TransformerBlock(nn.Module): + def __init__( self, dim: int, @@ -254,9 +253,7 @@ class MatrixGame3TransformerBlock(nn.Module): ) # 2. Cross-attention - self.attn2 = MatrixGame3CrossAttention( - dim, num_heads, qk_norm=qk_norm, eps=eps - ) + self.attn2 = MatrixGame3CrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, norm_type="layer", @@ -315,9 +312,7 @@ class MatrixGame3TransformerBlock(nn.Module): c_shift_msa, c_scale_msa, c_gate_msa, - ) = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk( - 6, dim=2 - ) + ) = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2) # batch_size, seq_len, 1, inner_dim shift_msa = shift_msa.squeeze(2) scale_msa = scale_msa.squeeze(2) @@ -339,9 +334,7 @@ class MatrixGame3TransformerBlock(nn.Module): assert shift_msa.dtype == torch.float32 # 1. Self-attention - norm_hidden_states = ( - self.norm1(hidden_states).float() * (1 + scale_msa) + shift_msa - ).to(orig_dtype) + norm_hidden_states = (self.norm1(hidden_states).float() * (1 + scale_msa) + shift_msa).to(orig_dtype) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -373,30 +366,42 @@ class MatrixGame3TransformerBlock(nn.Module): if pred_indices is None: pred_indices = torch.arange(grid_frames - memory_length, device=query.device, dtype=torch.long) - query_memory = _apply_rope_with_frame_indices( - query_memory, freqs, height=grid_height, width=grid_width, frame_indices=mem_indices - ) - key_memory = _apply_rope_with_frame_indices( - key_memory, freqs, height=grid_height, width=grid_width, frame_indices=mem_indices - ) - query_pred = _apply_rope_with_frame_indices( - query_pred, freqs, height=grid_height, width=grid_width, frame_indices=pred_indices - ) - key_pred = _apply_rope_with_frame_indices( - key_pred, freqs, height=grid_height, width=grid_width, frame_indices=pred_indices - ) + query_memory = _apply_rope_with_frame_indices(query_memory, + freqs, + height=grid_height, + width=grid_width, + frame_indices=mem_indices) + key_memory = _apply_rope_with_frame_indices(key_memory, + freqs, + height=grid_height, + width=grid_width, + frame_indices=mem_indices) + query_pred = _apply_rope_with_frame_indices(query_pred, + freqs, + height=grid_height, + width=grid_width, + frame_indices=pred_indices) + key_pred = _apply_rope_with_frame_indices(key_pred, + freqs, + height=grid_height, + width=grid_width, + frame_indices=pred_indices) query = torch.cat([query_memory, query_pred], dim=1) key = torch.cat([key_memory, key_pred], dim=1) else: pred_indices = predict_latent_idx if pred_indices is None: pred_indices = torch.arange(grid_frames, device=query.device, dtype=torch.long) - query = _apply_rope_with_frame_indices( - query, freqs, height=grid_height, width=grid_width, frame_indices=pred_indices - ) - key = _apply_rope_with_frame_indices( - key, freqs, height=grid_height, width=grid_width, frame_indices=pred_indices - ) + query = _apply_rope_with_frame_indices(query, + freqs, + height=grid_height, + width=grid_width, + frame_indices=pred_indices) + key = _apply_rope_with_frame_indices(key, + freqs, + height=grid_height, + width=grid_width, + frame_indices=pred_indices) attn_output, _ = self.attn1(query, key, value) attn_output = attn_output.flatten(2) @@ -407,8 +412,7 @@ class MatrixGame3TransformerBlock(nn.Module): hidden_states = hidden_states.float() + attn_output.float() * gate_msa.float() if self.use_memory and plucker_emb is not None: c2ws_hidden_states = self.cam_injector_layer2( - torch.nn.functional.silu(self.cam_injector_layer1(plucker_emb)) - ) + torch.nn.functional.silu(self.cam_injector_layer1(plucker_emb))) c2ws_hidden_states = c2ws_hidden_states + plucker_emb cam_scale = self.cam_scale_layer(c2ws_hidden_states) cam_shift = self.cam_shift_layer(c2ws_hidden_states) @@ -416,19 +420,14 @@ class MatrixGame3TransformerBlock(nn.Module): norm_hidden_states = self.self_attn_residual_norm.norm(hidden_states) # 2. Cross-attention - attn_output = self.attn2( - norm_hidden_states, context=encoder_hidden_states, context_lens=None - ) + attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None) # - use_memory / action path: x = norm3(x); x = x + cross_attn(x) # - plain path: x = x + cross_attn(norm3(x)) - cross_attn_residual_base = ( - norm_hidden_states if (mouse_cond is not None or self.use_memory) else hidden_states - ) + cross_attn_residual_base = (norm_hidden_states if + (mouse_cond is not None or self.use_memory) else hidden_states) hidden_states = cross_attn_residual_base + attn_output - norm_hidden_states = self.cross_attn_residual_norm.norm( - hidden_states.float() - ) + norm_hidden_states = self.cross_attn_residual_norm.norm(hidden_states.float()) norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa # ================= Action Module ================= @@ -446,9 +445,7 @@ class MatrixGame3TransformerBlock(nn.Module): keyboard_cond_memory=keyboard_cond_memory, ) - norm_hidden_states = self.cross_attn_residual_norm.norm( - hidden_states.float() - ) + norm_hidden_states = self.cross_attn_residual_norm.norm(hidden_states.float()) norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa # ================================================= @@ -469,16 +466,10 @@ class MatrixGame3WanModel(BaseDiT): _fsdp_shard_conditions = _DEFAULT_MATRIXGAME3_CONFIG._fsdp_shard_conditions _compile_conditions = _DEFAULT_MATRIXGAME3_CONFIG._compile_conditions - _supported_attention_backends = ( - _DEFAULT_MATRIXGAME3_CONFIG._supported_attention_backends - ) + _supported_attention_backends = (_DEFAULT_MATRIXGAME3_CONFIG._supported_attention_backends) param_names_mapping = _DEFAULT_MATRIXGAME3_CONFIG.param_names_mapping - reverse_param_names_mapping = ( - _DEFAULT_MATRIXGAME3_CONFIG.reverse_param_names_mapping - ) - lora_param_names_mapping = ( - _DEFAULT_MATRIXGAME3_CONFIG.lora_param_names_mapping - ) + reverse_param_names_mapping = (_DEFAULT_MATRIXGAME3_CONFIG.reverse_param_names_mapping) + lora_param_names_mapping = (_DEFAULT_MATRIXGAME3_CONFIG.lora_param_names_mapping) def __init__( self, @@ -531,24 +522,21 @@ class MatrixGame3WanModel(BaseDiT): self.action_config = getattr(config, "action_config", {}) # 3. Transformer blocks - self.blocks = nn.ModuleList( - [ - MatrixGame3TransformerBlock( - inner_dim, - config.ffn_dim, - config.num_attention_heads, - config.qk_norm, - config.cross_attn_norm, - config.eps, - self._supported_attention_backends, - prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", - action_config=self.action_config, - block_id=i, - use_memory=self.use_memory, - ) - for i in range(config.num_layers) - ] - ) + self.blocks = nn.ModuleList([ + MatrixGame3TransformerBlock( + inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.qk_norm, + config.cross_attn_norm, + config.eps, + self._supported_attention_backends, + prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", + action_config=self.action_config, + block_id=i, + use_memory=self.use_memory, + ) for i in range(config.num_layers) + ]) # 4. Output norm & projection self.norm_out = LayerNormScaleShift( @@ -559,12 +547,8 @@ class MatrixGame3WanModel(BaseDiT): dtype=torch.float32, compute_dtype=torch.float32, ) - self.proj_out = nn.Linear( - inner_dim, config.out_channels * math.prod(config.patch_size) - ) - self.scale_shift_table = nn.Parameter( - torch.randn(1, 2, inner_dim) / inner_dim**0.5 - ) + self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) self.gradient_checkpointing = False @@ -585,9 +569,7 @@ class MatrixGame3WanModel(BaseDiT): predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: - if encoder_hidden_states is not None and not isinstance( - encoder_hidden_states, torch.Tensor - ): + if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] memory_length = 0 @@ -617,23 +599,27 @@ class MatrixGame3WanModel(BaseDiT): predict_latent_idx = predict_latent_idx.to(device=device_idx, dtype=torch.long) elif isinstance(predict_latent_idx, tuple) and len(predict_latent_idx) == 2: predict_latent_idx = torch.arange( - predict_latent_idx[0], predict_latent_idx[1], - device=device_idx, dtype=torch.long, + predict_latent_idx[0], + predict_latent_idx[1], + device=device_idx, + dtype=torch.long, ) else: predict_latent_idx = torch.tensor( - list(predict_latent_idx), device=device_idx, dtype=torch.long, + list(predict_latent_idx), + device=device_idx, + dtype=torch.long, ) - if memory_latent_idx is None or ( - not torch.is_tensor(memory_latent_idx) and len(memory_latent_idx) == 0 - ): + if memory_latent_idx is None or (not torch.is_tensor(memory_latent_idx) and len(memory_latent_idx) == 0): memory_latent_idx = None elif torch.is_tensor(memory_latent_idx): memory_latent_idx = memory_latent_idx.to(device=device_idx, dtype=torch.long) else: memory_latent_idx = torch.tensor( - list(memory_latent_idx), device=device_idx, dtype=torch.long, + list(memory_latent_idx), + device=device_idx, + dtype=torch.long, ) head_dim = self.hidden_size // self.num_attention_heads @@ -646,10 +632,8 @@ class MatrixGame3WanModel(BaseDiT): ) max_seq_len = freqs.shape[1] if freqs.dim() == 3 else freqs.shape[0] if max_frame_index > max_seq_len or post_patch_height > max_seq_len or post_patch_width > max_seq_len: - raise ValueError( - f"got frames={max_frame_index}, height={post_patch_height}, width={post_patch_width}, " - f"max_seq_len={max_seq_len}" - ) + raise ValueError(f"got frames={max_frame_index}, height={post_patch_height}, width={post_patch_width}, " + f"max_seq_len={max_seq_len}") hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) @@ -668,20 +652,16 @@ class MatrixGame3WanModel(BaseDiT): target_frames = post_patch_num_frames * p_t target_height = post_patch_height * p_h target_width = post_patch_width * p_w - if c2ws_plucker_emb.shape[2] > target_frames or c2ws_plucker_emb.shape[3] > target_height or c2ws_plucker_emb.shape[4] > target_width: + if c2ws_plucker_emb.shape[2] > target_frames or c2ws_plucker_emb.shape[ + 3] > target_height or c2ws_plucker_emb.shape[4] > target_width: c2ws_plucker_emb = c2ws_plucker_emb[:, :, :target_frames, :target_height, :target_width] if c2ws_plucker_emb.shape[2:] != (target_frames, target_height, target_width): - raise ValueError( - "Camera conditioning must align with the patchified latent grid, " - f"got camera shape {tuple(c2ws_plucker_emb.shape[2:])} vs target " - f"{(target_frames, target_height, target_width)}" - ) - plucker_emb = self.camera_patch_embedding( - c2ws_plucker_emb.to(hidden_states.dtype) - ) + raise ValueError("Camera conditioning must align with the patchified latent grid, " + f"got camera shape {tuple(c2ws_plucker_emb.shape[2:])} vs target " + f"{(target_frames, target_height, target_width)}") + plucker_emb = self.camera_patch_embedding(c2ws_plucker_emb.to(hidden_states.dtype)) plucker_hidden = self.c2ws_hidden_states_layer2( - torch.nn.functional.silu(self.c2ws_hidden_states_layer1(plucker_emb)) - ) + torch.nn.functional.silu(self.c2ws_hidden_states_layer1(plucker_emb))) plucker_emb = plucker_emb + plucker_hidden timestep_tokens = timestep @@ -689,20 +669,15 @@ class MatrixGame3WanModel(BaseDiT): timestep_tokens = timestep_tokens.unsqueeze(0) if timestep_tokens.dim() == 1: timestep_tokens = timestep_tokens.unsqueeze(1).repeat( - 1, post_patch_num_frames * post_patch_height * post_patch_width - ) + 1, post_patch_num_frames * post_patch_height * post_patch_width) elif timestep_tokens.dim() == 2 and timestep_tokens.shape[1] == num_frames: - timestep_tokens = ( - timestep_tokens[:, :, None, None] - .repeat(1, 1, post_patch_height, post_patch_width) - .reshape(timestep_tokens.shape[0], -1) - ) + timestep_tokens = (timestep_tokens[:, :, None, + None].repeat(1, 1, post_patch_height, + post_patch_width).reshape(timestep_tokens.shape[0], -1)) if memory_length > 0: if timestep_memory is None: raise ValueError("MatrixGame3 requires timestep_memory when x_memory is provided") - timestep_tokens = torch.cat( - [timestep_memory.to(timestep_tokens.dtype), timestep_tokens], dim=1 - ) + timestep_tokens = torch.cat([timestep_memory.to(timestep_tokens.dtype), timestep_tokens], dim=1) temb, timestep_proj, encoder_hidden_states = self.condition_embedder( timestep_tokens.flatten(), @@ -717,13 +692,9 @@ class MatrixGame3WanModel(BaseDiT): encoder_hidden_states = encoder_hidden_states[0] elif encoder_hidden_states.ndim == 2: encoder_hidden_states = encoder_hidden_states.unsqueeze(0) - if ( - encoder_hidden_states.ndim == 3 - and encoder_hidden_states.shape[-1] == self.text_embedding[0].in_features - ): - encoder_hidden_states = self.text_embedding( - encoder_hidden_states.to(hidden_states.dtype) - ) + if (encoder_hidden_states.ndim == 3 + and encoder_hidden_states.shape[-1] == self.text_embedding[0].in_features): + encoder_hidden_states = self.text_embedding(encoder_hidden_states.to(hidden_states.dtype)) # [F, H, W] for the ActionModule grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width) @@ -766,15 +737,11 @@ class MatrixGame3WanModel(BaseDiT): # Output if temb.dim() == 3: - shift, scale = ( - self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2) - ).chunk(2, dim=2) + shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2) shift = shift.squeeze(2) scale = scale.squeeze(2) else: - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk( - 2, dim=1 - ) + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) diff --git a/fastvideo/models/dits/matrixgame3/utils.py b/fastvideo/models/dits/matrixgame3/utils.py index f51d018f3..55232ea43 100644 --- a/fastvideo/models/dits/matrixgame3/utils.py +++ b/fastvideo/models/dits/matrixgame3/utils.py @@ -11,7 +11,6 @@ from fastvideo.models.dits.lingbotworld.cam_utils import ( interpolate_camera_poses, ) - WSAD_OFFSET = 12.35 DIAGONAL_OFFSET = 8.73 MOUSE_PITCH_SENSITIVITY = 15.0 @@ -236,7 +235,8 @@ def select_memory_idx_fov( grid_x.reshape(-1) * grid_z.reshape(-1) * (width / (2 * fx)), grid_y.reshape(-1) * grid_z.reshape(-1) * (height / (2 * fy)), grid_z.reshape(-1), - ], dim=0) + ], + dim=0) selected_index: list[int] = [] selected_confidence: list[float] = [] @@ -244,7 +244,8 @@ def select_memory_idx_fov( base_pose = extrinsics_tensor[frame_idx] points_world = base_pose[:3, :3] @ points_cam_base + base_pose[:3, 3:4] points_world_batched = points_world.unsqueeze(0) - points_in_cands = torch.bmm(r_cand_inv, points_world_batched.expand(len(candidate_indices), -1, -1)) + t_cand_inv + points_in_cands = torch.bmm(r_cand_inv, points_world_batched.expand(len(candidate_indices), -1, + -1)) + t_cand_inv x = points_in_cands[:, 0, :] y = points_in_cands[:, 1, :] @@ -295,9 +296,7 @@ def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = No actions_single_camera = [] keyboard_idx = {"still": 0, "forward": 1, "back": 2, "left": 3, "right": 4, "a": 5, "d": 6} - actions_to_test = ( - actions_double_action * 5 + actions_single_camera * 5 + actions_single_action * 5 - ) + actions_to_test = (actions_double_action * 5 + actions_single_camera * 5 + actions_single_action * 5) for action in (actions_single_action + actions_double_action): for camera in actions_single_camera: actions_to_test.append(f"{action}_{camera}") diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index b06f36f83..44c8a540f 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -23,7 +23,6 @@ from fastvideo.layers.visual_embedding import Timesteps from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum - MINIMAX_H3_MODALITY_NUM = 3 _CFG = MiniMaxH3Config() @@ -33,9 +32,8 @@ class MiniMaxH3RotaryPosEmbed(nn.Module): def __init__(self, rope_freq_dim: int, rope_theta: float) -> None: super().__init__() - inv_freq = 1.0 / ( - rope_theta ** (torch.arange(0, 2 * rope_freq_dim, 2, dtype=torch.float32) / (2 * rope_freq_dim)) - ) + inv_freq = 1.0 / (rope_theta**(torch.arange(0, 2 * rope_freq_dim, 2, dtype=torch.float32) / + (2 * rope_freq_dim))) self.register_buffer("inv_freq", inv_freq, persistent=False) def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -355,9 +353,7 @@ class MiniMaxH3AdaLayerNormOut(nn.Module): shift_scale, _ = self.linear(F.silu(temb).to(self.linear.weight.dtype)) shift, scale = shift_scale.chunk(2, dim=-1) hidden_states = self.norm(hidden_states) - return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( - 0, timestep_indices - ) + return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select(0, timestep_indices) class MiniMaxH3TransformerBlock(nn.Module): @@ -414,16 +410,14 @@ class MiniMaxH3TransformerBlock(nn.Module): residual = hidden_states norm_hidden_states = self.norm1(hidden_states) norm_hidden_states = norm_hidden_states * ( - 1.0 + scale_msa.index_select(0, adaln_indices) - ) + shift_msa.index_select(0, adaln_indices) + 1.0 + scale_msa.index_select(0, adaln_indices)) + shift_msa.index_select(0, adaln_indices) attention_output = self.attn(norm_hidden_states, rotary_emb, original_seq_len) hidden_states = residual + gate_msa.index_select(0, adaln_indices) * attention_output residual = hidden_states norm_hidden_states = self.norm2(hidden_states) norm_hidden_states = norm_hidden_states * ( - 1.0 + scale_mlp.index_select(0, adaln_indices) - ) + shift_mlp.index_select(0, adaln_indices) + 1.0 + scale_mlp.index_select(0, adaln_indices)) + shift_mlp.index_select(0, adaln_indices) feed_forward_output = self.ff(norm_hidden_states) return residual + gate_mlp.index_select(0, adaln_indices) * feed_forward_output @@ -467,10 +461,8 @@ class MiniMaxH3Transformer3DModel(BaseDiT): arch = config.arch_config sp_world_size = get_sp_world_size() if model_parallel_is_initialized() else 1 if arch.num_attention_heads % sp_world_size: - raise ValueError( - f"MiniMax H3 attention heads ({arch.num_attention_heads}) must be divisible by " - f"sequence parallel size ({sp_world_size})." - ) + raise ValueError(f"MiniMax H3 attention heads ({arch.num_attention_heads}) must be divisible by " + f"sequence parallel size ({sp_world_size}).") self.hidden_size = arch.hidden_size self.num_attention_heads = arch.num_attention_heads @@ -576,13 +568,9 @@ class MiniMaxH3Transformer3DModel(BaseDiT): del dtype if self.rope.inv_freq.is_meta or self.rope.inv_freq.device != device: arch = self.config.arch_config - inv_freq = 1.0 / ( - arch.rope_theta - ** ( - torch.arange(0, 2 * arch.rope_freq_dim, 2, device=device, dtype=torch.float32) - / (2 * arch.rope_freq_dim) - ) - ) + inv_freq = 1.0 / (arch.rope_theta + **(torch.arange(0, 2 * arch.rope_freq_dim, 2, device=device, dtype=torch.float32) / + (2 * arch.rope_freq_dim))) self.rope._buffers["inv_freq"] = inv_freq def forward( @@ -602,7 +590,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT): if position_ids.ndim != 2 or position_ids.shape[-1] != 3: raise ValueError(f"position_ids must have shape (seq_len, 3), got {tuple(position_ids.shape)}.") sequence_length = position_ids.shape[0] - if token_tags.shape != (sequence_length,) or timestep_indices.shape != (sequence_length,): + if token_tags.shape != (sequence_length, ) or timestep_indices.shape != (sequence_length, ): raise ValueError("token_tags and timestep_indices must both match the packed sequence length.") if hidden_states.shape[1] != video_indices.numel(): raise ValueError("hidden_states row count must match video_indices.") @@ -628,9 +616,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT): dim=1, ) - packed_hidden_states = text_embeds.new_zeros( - (text_embeds.shape[0], sequence_length, text_embeds.shape[-1]) - ) + packed_hidden_states = text_embeds.new_zeros((text_embeds.shape[0], sequence_length, text_embeds.shape[-1])) packed_hidden_states = packed_hidden_states.index_copy(1, text_indices, text_embeds) packed_hidden_states = packed_hidden_states.index_copy(1, video_indices, video_embeds.to(text_embeds.dtype)) packed_hidden_states = packed_hidden_states.index_copy(1, audio_indices, audio_embeds.to(text_embeds.dtype)) diff --git a/fastvideo/models/dits/sd3.py b/fastvideo/models/dits/sd3.py index 6f36cb17e..6b6f4bf81 100644 --- a/fastvideo/models/dits/sd3.py +++ b/fastvideo/models/dits/sd3.py @@ -33,10 +33,8 @@ def _chunked_feed_forward( chunk_size: int, ) -> torch.Tensor: if hidden_states.shape[chunk_dim] % chunk_size != 0: - raise ValueError( - f"hidden_states.shape[{chunk_dim}] ({hidden_states.shape[chunk_dim]})" - f" must be divisible by chunk_size ({chunk_size})" - ) + raise ValueError(f"hidden_states.shape[{chunk_dim}] ({hidden_states.shape[chunk_dim]})" + f" must be divisible by chunk_size ({chunk_size})") num_chunks = hidden_states.shape[chunk_dim] // chunk_size return torch.cat( [ff(h) for h in hidden_states.chunk(num_chunks, dim=chunk_dim)], @@ -89,16 +87,10 @@ def _get_2d_sincos_pos_embed( if isinstance(grid_size, int): grid_size = (grid_size, grid_size) - grid_h = ( - torch.arange(grid_size[0], device=device, dtype=torch.float32) - / (grid_size[0] / base_size) - / interpolation_scale - ) - grid_w = ( - torch.arange(grid_size[1], device=device, dtype=torch.float32) - / (grid_size[1] / base_size) - / interpolation_scale - ) + grid_h = (torch.arange(grid_size[0], device=device, dtype=torch.float32) / (grid_size[0] / base_size) / + interpolation_scale) + grid_w = (torch.arange(grid_size[1], device=device, dtype=torch.float32) / (grid_size[1] / base_size) / + interpolation_scale) grid = torch.meshgrid(grid_w, grid_h, indexing="xy") grid = torch.stack(grid, dim=0) grid = grid.reshape([2, 1, grid_size[1], grid_size[0]]) @@ -137,8 +129,7 @@ class SD3PatchEmbed(nn.Module): bias=bias, ) if layer_norm: - self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, - eps=1e-6) + self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) else: self.norm = None @@ -148,7 +139,7 @@ class SD3PatchEmbed(nn.Module): self.base_size = height // patch_size self.interpolation_scale = interpolation_scale - grid_size = pos_embed_max_size or int(num_patches ** 0.5) + grid_size = pos_embed_max_size or int(num_patches**0.5) if pos_embed_type is None: self.pos_embed = None @@ -176,15 +167,11 @@ class SD3PatchEmbed(nn.Module): width = width // self.patch_size if height > self.pos_embed_max_size: - raise ValueError( - f"Height ({height}) cannot be greater than pos_embed_max_size " - f"({self.pos_embed_max_size})" - ) + raise ValueError(f"Height ({height}) cannot be greater than pos_embed_max_size " + f"({self.pos_embed_max_size})") if width > self.pos_embed_max_size: - raise ValueError( - f"Width ({width}) cannot be greater than pos_embed_max_size " - f"({self.pos_embed_max_size})" - ) + raise ValueError(f"Width ({width}) cannot be greater than pos_embed_max_size " + f"({self.pos_embed_max_size})") top = (self.pos_embed_max_size - height) // 2 left = (self.pos_embed_max_size - width) // 2 @@ -243,11 +230,9 @@ class SD3TimestepEmbedding(nn.Module): ) -> None: super().__init__() - self.linear_1 = ReplicatedLinear(in_channels, time_embed_dim, - bias=True) + self.linear_1 = ReplicatedLinear(in_channels, time_embed_dim, bias=True) self.act = get_act_fn(act_fn) - self.linear_2 = ReplicatedLinear(time_embed_dim, time_embed_dim, - bias=True) + self.linear_2 = ReplicatedLinear(time_embed_dim, time_embed_dim, bias=True) def forward(self, sample: torch.Tensor) -> torch.Tensor: sample, _ = self.linear_1(sample) @@ -308,8 +293,7 @@ class CombinedTimestepTextProjEmbeddings(nn.Module): pooled_projection: torch.Tensor, ) -> torch.Tensor: timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder( - timesteps_proj.to(dtype=pooled_projection.dtype)) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype)) pooled_projections = self.text_embedder(pooled_projection) return timesteps_emb + pooled_projections @@ -325,16 +309,11 @@ class SD35AdaLayerNormZeroX(nn.Module): super().__init__() self.silu = nn.SiLU() - self.linear = ReplicatedLinear(embedding_dim, 9 * embedding_dim, - bias=bias) + self.linear = ReplicatedLinear(embedding_dim, 9 * embedding_dim, bias=bias) if norm_type != "layer_norm": - raise ValueError( - f"Unsupported norm_type ({norm_type}); expected layer_norm" - ) - self.norm = nn.LayerNorm(embedding_dim, - elementwise_affine=False, - eps=1e-6) + raise ValueError(f"Unsupported norm_type ({norm_type}); expected layer_norm") + self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) def forward( self, @@ -342,15 +321,12 @@ class SD35AdaLayerNormZeroX(nn.Module): emb: torch.Tensor, ) -> tuple[torch.Tensor, ...]: emb, _ = self.linear(self.silu(emb)) - (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, - shift_msa2, scale_msa2, + (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, shift_msa2, scale_msa2, gate_msa2) = emb.chunk(9, dim=1) norm_hidden_states = self.norm(hidden_states) - hidden_states = norm_hidden_states * ( - 1 + scale_msa[:, None]) + shift_msa[:, None] - norm_hidden_states2 = norm_hidden_states * ( - 1 + scale_msa2[:, None]) + shift_msa2[:, None] + hidden_states = norm_hidden_states * (1 + scale_msa[:, None]) + shift_msa[:, None] + norm_hidden_states2 = norm_hidden_states * (1 + scale_msa2[:, None]) + shift_msa2[:, None] return ( hidden_states, gate_msa, @@ -373,26 +349,19 @@ class SD3AdaLayerNormZero(nn.Module): super().__init__() self.silu = nn.SiLU() - self.linear = ReplicatedLinear(embedding_dim, 6 * embedding_dim, - bias=bias) + self.linear = ReplicatedLinear(embedding_dim, 6 * embedding_dim, bias=bias) if norm_type != "layer_norm": - raise ValueError( - f"Unsupported norm_type ({norm_type}); expected layer_norm" - ) - self.norm = nn.LayerNorm(embedding_dim, - elementwise_affine=False, - eps=1e-6) + raise ValueError(f"Unsupported norm_type ({norm_type}); expected layer_norm") + self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) def forward( self, x: torch.Tensor, emb: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, - torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: emb, _ = self.linear(self.silu(emb)) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - emb.chunk(6, dim=1)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (emb.chunk(6, dim=1)) x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] return x, gate_msa, shift_mlp, scale_mlp, gate_mlp @@ -536,8 +505,7 @@ class SD3Attention(nn.Module): context_pre_only: bool | None = None, qk_norm: str | None = None, eps: float = 1e-6, - supported_attention_backends: tuple[AttentionBackendEnum, - ...] | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, ) -> None: super().__init__() @@ -585,8 +553,7 @@ class SD3Attention(nn.Module): ]) if context_pre_only is not None and not context_pre_only: - self.to_add_out = ReplicatedLinear(self.inner_dim, out_dim, - bias=True) + self.to_add_out = ReplicatedLinear(self.inner_dim, out_dim, bias=True) else: self.to_add_out = None @@ -619,12 +586,9 @@ class SD3Attention(nn.Module): key = self.norm_k(key) if encoder_hidden_states is not None: - if (self.add_q_proj is None or self.add_k_proj is None - or self.add_v_proj is None): - raise RuntimeError( - "encoder_hidden_states provided but attention does not have" - " added projections" - ) + if (self.add_q_proj is None or self.add_k_proj is None or self.add_v_proj is None): + raise RuntimeError("encoder_hidden_states provided but attention does not have" + " added projections") encoder_query, _ = self.add_q_proj(encoder_hidden_states) encoder_key, _ = self.add_k_proj(encoder_hidden_states) @@ -659,8 +623,7 @@ class SD3Attention(nn.Module): value = torch.cat([value, encoder_value], dim=1) attn_output, _ = self.attn(query, key, value) - attn_output = attn_output.reshape(batch_size, -1, - self.heads * self.head_dim) + attn_output = attn_output.reshape(batch_size, -1, self.heads * self.head_dim) if encoder_hidden_states is not None: hidden_out = attn_output[:, :residual_seq_len] @@ -689,8 +652,7 @@ class SD3JointTransformerBlock(nn.Module): context_pre_only: bool = False, qk_norm: str | None = None, use_dual_attention: bool = False, - supported_attention_backends: tuple[AttentionBackendEnum, - ...] | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, ) -> None: super().__init__() @@ -749,9 +711,7 @@ class SD3JointTransformerBlock(nn.Module): ) if not context_pre_only: - self.norm2_context = nn.LayerNorm(dim, - elementwise_affine=False, - eps=1e-6) + self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) self.ff_context = SD3FeedForward( dim=dim, dim_out=dim, @@ -823,8 +783,7 @@ class SD3JointTransformerBlock(nn.Module): hidden_states = hidden_states + gate_msa2.unsqueeze(1) * attn_output2 norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * ( - 1 + scale_mlp[:, None]) + shift_mlp[:, None] + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] if self._chunk_size is not None: ff_output = _chunked_feed_forward( @@ -849,8 +808,7 @@ class SD3JointTransformerBlock(nn.Module): encoder_hidden_states = encoder_hidden_states + context_attn_output norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) - norm_encoder_hidden_states = norm_encoder_hidden_states * ( - 1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] if self._chunk_size is not None: context_ff_output = _chunked_feed_forward( @@ -862,8 +820,7 @@ class SD3JointTransformerBlock(nn.Module): else: context_ff_output = self.ff_context(norm_encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states + ( - c_gate_mlp.unsqueeze(1) * context_ff_output) + encoder_hidden_states = encoder_hidden_states + (c_gate_mlp.unsqueeze(1) * context_ff_output) return encoder_hidden_states, hidden_states @@ -913,12 +870,10 @@ class SD3Transformer2DModel(BaseDiT): embedding_dim=self.inner_dim, pooled_projection_dim=arch.pooled_projection_dim, ) - self.context_embedder = ReplicatedLinear(arch.joint_attention_dim, - arch.caption_projection_dim) + self.context_embedder = ReplicatedLinear(arch.joint_attention_dim, arch.caption_projection_dim) dual_layers = getattr(arch, "dual_attention_layers", ()) - dual_layers = tuple(dual_layers) if isinstance(dual_layers, - list) else dual_layers + dual_layers = tuple(dual_layers) if isinstance(dual_layers, list) else dual_layers self.transformer_blocks = nn.ModuleList([ SD3JointTransformerBlock( @@ -959,8 +914,7 @@ class SD3Transformer2DModel(BaseDiT): chunk_size = chunk_size or 1 - def fn_recursive_feed_forward(module: torch.nn.Module, chunk: int, - chunk_dim: int): + def fn_recursive_feed_forward(module: torch.nn.Module, chunk: int, chunk_dim: int): if hasattr(module, "set_chunk_feed_forward"): module.set_chunk_feed_forward(chunk_size=chunk, dim=chunk_dim) for child in module.children(): @@ -1023,12 +977,10 @@ class SD3Transformer2DModel(BaseDiT): hidden_states = self.pos_embed(hidden_states) temb = self.time_text_embed(timestep, pooled_projections) - encoder_hidden_states, _ = self.context_embedder( - encoder_hidden_states) + encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states) for index_block, block in enumerate(self.transformer_blocks): - is_skip = bool( - skip_layers is not None and index_block in skip_layers) + is_skip = bool(skip_layers is not None and index_block in skip_layers) if is_skip: continue @@ -1039,12 +991,9 @@ class SD3Transformer2DModel(BaseDiT): joint_attention_kwargs=joint_attention_kwargs, ) - if (block_controlnet_hidden_states is not None - and not block.context_pre_only): - interval_control = len(self.transformer_blocks) / len( - block_controlnet_hidden_states) - hidden_states = hidden_states + block_controlnet_hidden_states[ - int(index_block / interval_control)] + if (block_controlnet_hidden_states is not None and not block.context_pre_only): + interval_control = len(self.transformer_blocks) / len(block_controlnet_hidden_states) + hidden_states = hidden_states + block_controlnet_hidden_states[int(index_block / interval_control)] hidden_states = self.norm_out(hidden_states, temb) hidden_states, _ = self.proj_out(hidden_states) diff --git a/fastvideo/models/dits/stable_audio.py b/fastvideo/models/dits/stable_audio.py index 693f2cad5..ee1f5808b 100644 --- a/fastvideo/models/dits/stable_audio.py +++ b/fastvideo/models/dits/stable_audio.py @@ -124,8 +124,12 @@ class FeedForward(nn.Module): class Attention(nn.Module): - def __init__(self, dim: int, dim_heads: int = 64, dim_context: int | None = None, - zero_init_output: bool = True, qk_norm: str | None = None) -> None: + def __init__(self, + dim: int, + dim_heads: int = 64, + dim_context: int | None = None, + zero_init_output: bool = True, + qk_norm: str | None = None) -> None: super().__init__() self.dim = dim self.dim_heads = dim_heads @@ -154,11 +158,15 @@ class Attention(nn.Module): else: raise ValueError(f"Unsupported qk_norm={qk_norm!r}; expected 'ln' or None.") - self.attn = LocalAttention(num_heads=self.num_heads, head_size=dim_heads, - num_kv_heads=self.kv_heads, causal=False, + self.attn = LocalAttention(num_heads=self.num_heads, + head_size=dim_heads, + num_kv_heads=self.kv_heads, + causal=False, supported_attention_backends=_SUPPORTED_BACKENDS) - def forward(self, x: torch.Tensor, context: torch.Tensor | None = None, + def forward(self, + x: torch.Tensor, + context: torch.Tensor | None = None, rotary_pos_emb: tuple[torch.Tensor, float] | None = None) -> torch.Tensor: h, kv_h, has_context = self.num_heads, self.kv_heads, context is not None kv_input = context if has_context else x @@ -201,26 +209,35 @@ class Attention(nn.Module): class TransformerBlock(nn.Module): - def __init__(self, dim: int, dim_heads: int = 64, cross_attend: bool = False, - dim_context: int | None = None, zero_init_branch_outputs: bool = True, + def __init__(self, + dim: int, + dim_heads: int = 64, + cross_attend: bool = False, + dim_context: int | None = None, + zero_init_branch_outputs: bool = True, qk_norm: str | None = None) -> None: super().__init__() self.dim = dim self.dim_heads = min(dim_heads, dim) self.cross_attend = cross_attend self.pre_norm = FP32LayerNorm(dim, elementwise_affine=True) - self.self_attn = Attention(dim, dim_heads=self.dim_heads, + self.self_attn = Attention(dim, + dim_heads=self.dim_heads, zero_init_output=zero_init_branch_outputs, qk_norm=qk_norm) if cross_attend: self.cross_attend_norm = FP32LayerNorm(dim, elementwise_affine=True) - self.cross_attn = Attention(dim, dim_heads=self.dim_heads, dim_context=dim_context, + self.cross_attn = Attention(dim, + dim_heads=self.dim_heads, + dim_context=dim_context, zero_init_output=zero_init_branch_outputs, qk_norm=qk_norm) self.ff_norm = FP32LayerNorm(dim, elementwise_affine=True) self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs) - def forward(self, x: torch.Tensor, context: torch.Tensor | None = None, + def forward(self, + x: torch.Tensor, + context: torch.Tensor | None = None, rotary_pos_emb: tuple[torch.Tensor, float] | None = None) -> torch.Tensor: x = x + self.self_attn(self.pre_norm(x), rotary_pos_emb=rotary_pos_emb) if context is not None and self.cross_attend: @@ -231,26 +248,35 @@ class TransformerBlock(nn.Module): class ContinuousTransformer(nn.Module): - def __init__(self, dim: int, depth: int, *, dim_heads: int = 64, dim_in: int | None = None, - dim_out: int | None = None, cross_attend: bool = False, - cond_token_dim: int | None = None, zero_init_branch_outputs: bool = True, + def __init__(self, + dim: int, + depth: int, + *, + dim_heads: int = 64, + dim_in: int | None = None, + dim_out: int | None = None, + cross_attend: bool = False, + cond_token_dim: int | None = None, + zero_init_branch_outputs: bool = True, qk_norm: str | None = None) -> None: super().__init__() self.dim = dim self.depth = depth - self.project_in = (ReplicatedLinear(dim_in, dim, bias=False) if dim_in is not None - else nn.Identity()) - self.project_out = (ReplicatedLinear(dim, dim_out, bias=False) if dim_out is not None - else nn.Identity()) + self.project_in = (ReplicatedLinear(dim_in, dim, bias=False) if dim_in is not None else nn.Identity()) + self.project_out = (ReplicatedLinear(dim, dim_out, bias=False) if dim_out is not None else nn.Identity()) self.rotary_pos_emb = RotaryEmbedding(max(dim_heads // 2, 32)) self.layers = nn.ModuleList([ - TransformerBlock(dim, dim_heads=dim_heads, cross_attend=cross_attend, + TransformerBlock(dim, + dim_heads=dim_heads, + cross_attend=cross_attend, dim_context=cond_token_dim, zero_init_branch_outputs=zero_init_branch_outputs, qk_norm=qk_norm) for _ in range(depth) ]) - def forward(self, x: torch.Tensor, prepend_embeds: torch.Tensor | None = None, + def forward(self, + x: torch.Tensor, + prepend_embeds: torch.Tensor | None = None, context: torch.Tensor | None = None) -> torch.Tensor: if isinstance(self.project_in, ReplicatedLinear): x, _ = self.project_in(x) @@ -273,8 +299,7 @@ class StableAudioDiT(BaseDiT): param_names_mapping = _DEFAULT_CONFIG.arch_config.param_names_mapping reverse_param_names_mapping: dict = {} - def __init__(self, config: StableAudioConfig | None = None, - hf_config: dict[str, Any] | None = None) -> None: + def __init__(self, config: StableAudioConfig | None = None, hf_config: dict[str, Any] | None = None) -> None: if config is None: config = StableAudioConfig() super().__init__(config=config, hf_config=hf_config or {}) @@ -317,8 +342,13 @@ class StableAudioDiT(BaseDiT): ) self.transformer = ContinuousTransformer( - dim=embed_dim, depth=depth, dim_heads=embed_dim // num_heads, dim_in=io_channels, - dim_out=io_channels, cross_attend=True, cond_token_dim=cond_embed_dim, + dim=embed_dim, + depth=depth, + dim_heads=embed_dim // num_heads, + dim_in=io_channels, + dim_out=io_channels, + cross_attend=True, + cond_token_dim=cond_embed_dim, qk_norm=qk_norm, ) @@ -364,7 +394,8 @@ class StableAudioDiT(BaseDiT): return self.postprocess_conv(out) + out @classmethod - def from_official_state_dict(cls, state_dict: dict[str, torch.Tensor], + def from_official_state_dict(cls, + state_dict: dict[str, torch.Tensor], prefix: str = "model.model.") -> "StableAudioDiT": """Load from a raw `stable_audio_tools` monolithic state dict. Kept for tests / older checkpoints; production loads go through @@ -381,8 +412,7 @@ class StableAudioDiT(BaseDiT): remapped[new_key] = v missing, unexpected = model.load_state_dict(remapped, strict=True) if missing or unexpected: - raise RuntimeError( - f"StableAudioDiT load mismatch — missing={missing[:5]} unexpected={unexpected[:5]}") + raise RuntimeError(f"StableAudioDiT load mismatch — missing={missing[:5]} unexpected={unexpected[:5]}") return model diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 6a153c2ee..4c971ca9e 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -8,22 +8,18 @@ import torch import torch.nn as nn import fastvideo.envs as envs -from fastvideo.attention import (DistributedAttention, DistributedAttention_VSA, - LocalAttention) +from fastvideo.attention import (DistributedAttention, DistributedAttention_VSA, LocalAttention) from fastvideo.configs.models.dits import WanVideoConfig -from fastvideo.distributed.communication_op import ( - sequence_model_parallel_all_gather_with_unpad, - sequence_model_parallel_shard) -from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, - RMSNorm, ScaleResidual, +from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad, + sequence_model_parallel_shard) +from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual, ScaleResidualLayerNormScaleShift) from fastvideo.layers.linear import ReplicatedLinear # from torch.nn import RMSNorm # TODO: RMSNorm .... from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import get_rotary_pos_embed -from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, - TimestepEmbedder) +from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, TimestepEmbedder) from fastvideo.logger import init_logger from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum, current_platform @@ -43,8 +39,7 @@ class WanImageEmbedding(torch.nn.Module): self.ff = MLP(in_features, in_features, out_features, act_type="gelu") self.norm2 = FP32LayerNorm(out_features) - def forward(self, - encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: + def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: dtype = encoder_hidden_states_image.dtype hidden_states = self.norm1(encoder_hidden_states_image) hidden_states = self.ff(hidden_states) @@ -68,15 +63,9 @@ class WanTimeTextImageEmbedding(nn.Module): ): super().__init__() - self.time_embedder = TimestepEmbedder( - dim, frequency_embedding_size=time_freq_dim, act_layer="silu") - self.time_modulation = ModulateProjection(dim, - factor=6, - act_layer="silu") - self.text_embedder = MLP(text_embed_dim, - dim, - dim, - bias=True, + self.time_embedder = TimestepEmbedder(dim, frequency_embedding_size=time_freq_dim, act_layer="silu") + self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") + self.text_embedder = MLP(text_embed_dim, dim, dim, bias=True, act_type="gelu_pytorch_tanh") if text_embed_dim > 0 else None self.image_embedder = None @@ -91,13 +80,11 @@ class WanTimeTextImageEmbedding(nn.Module): self._r_embedder_deltatime_type = r_embedder_deltatime_type if self._r_embedder_enabled: if r_embedder_fusion not in ("additive", "gated"): - raise ValueError( - "r_embedder_fusion must be one of {additive, gated}, " - f"got {r_embedder_fusion!r}") + raise ValueError("r_embedder_fusion must be one of {additive, gated}, " + f"got {r_embedder_fusion!r}") if r_embedder_deltatime_type not in ("r", "t-r"): - raise ValueError( - "r_embedder_deltatime_type must be one of {r, t-r}, " - f"got {r_embedder_deltatime_type!r}") + raise ValueError("r_embedder_deltatime_type must be one of {r, t-r}, " + f"got {r_embedder_deltatime_type!r}") # Deep-copy preserves identical initialization with time_embedder, # matching AnyFlow reference setup_flowmap_model() behavior. self.delta_embedder = copy.deepcopy(self.time_embedder) @@ -139,11 +126,12 @@ class WanTimeTextImageEmbedding(nn.Module): if self.text_embedder is not None: encoder_hidden_states = self.text_embedder(encoder_hidden_states) else: - encoder_hidden_states = torch.zeros((timestep.shape[0], 0, temb.shape[-1]), device=temb.device, dtype=temb.dtype) + encoder_hidden_states = torch.zeros((timestep.shape[0], 0, temb.shape[-1]), + device=temb.device, + dtype=temb.dtype) if encoder_hidden_states_image is not None: assert self.image_embedder is not None - encoder_hidden_states_image = self.image_embedder( - encoder_hidden_states_image) + encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image) return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image @@ -178,17 +166,15 @@ class WanSelfAttention(nn.Module): self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() # Scaled dot product attention - self.attn = LocalAttention( - num_heads=num_heads, - head_size=self.head_dim, - dropout_rate=0, - softmax_scale=None, - causal=False, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA)) + self.attn = LocalAttention(num_heads=num_heads, + head_size=self.head_dim, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA)) - def forward(self, x: torch.Tensor, context: torch.Tensor, - context_lens: int): + def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: int): r""" Args: x(Tensor): Shape [B, L, num_heads, C / num_heads] @@ -250,8 +236,14 @@ class WanI2VCrossAttention(WanSelfAttention): quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: - super().__init__(dim, num_heads, window_size, qk_norm, eps, - supported_attention_backends, quant_config=quant_config, prefix=prefix) + super().__init__(dim, + num_heads, + window_size, + qk_norm, + eps, + supported_attention_backends, + quant_config=quant_config, + prefix=prefix) self.add_k_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_k_proj") self.add_v_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_v_proj") @@ -273,8 +265,7 @@ class WanI2VCrossAttention(WanSelfAttention): q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) v = self.to_v(context)[0].view(b, -1, n, d) - k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view( - b, -1, n, d) + k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view(b, -1, n, d) v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) img_x = self.attn(q, k_img, v_img) # compute attention @@ -311,12 +302,11 @@ class WanTransformerBlock(nn.Module): self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v") self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out") - self.attn1 = DistributedAttention( - num_heads=num_heads, - head_size=dim // num_heads, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.attn1") + self.attn1 = DistributedAttention(num_heads=num_heads, + head_size=dim // num_heads, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn1") self.hidden_dim = dim self.num_attention_heads = num_heads dim_head = dim // num_heads @@ -331,13 +321,12 @@ class WanTransformerBlock(nn.Module): print("QK Norm type not supported") raise Exception assert cross_attn_norm is True - self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( - dim, - norm_type="layer", - eps=eps, - elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32) + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32, + compute_dtype=torch.float32) # 2. Cross-attention if added_kv_proj_dim is not None: @@ -356,13 +345,12 @@ class WanTransformerBlock(nn.Module): eps=eps, quant_config=quant_config, prefix=f"{prefix}.attn2") - self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( - dim, - norm_type="layer", - eps=eps, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn") @@ -387,8 +375,7 @@ class WanTransformerBlock(nn.Module): if temb.dim() == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - self.scale_shift_table.unsqueeze(0) + temb.float() - ).chunk(6, dim=2) + self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2) # batch_size, seq_len, 1, inner_dim shift_msa = shift_msa.squeeze(2) scale_msa = scale_msa.squeeze(2) @@ -399,13 +386,11 @@ class WanTransformerBlock(nn.Module): else: # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) e = self.scale_shift_table + temb.float() - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( - 6, dim=1) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1) assert shift_msa.dtype == torch.float32 # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()) * - (1 + scale_msa) + shift_msa).to(orig_dtype) + norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -431,19 +416,15 @@ class WanTransformerBlock(nn.Module): attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.tensor([0], device=hidden_states.device) - norm_hidden_states, hidden_states = self.self_attn_residual_norm( - hidden_states, attn_output, gate_msa, null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, + null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) # 2. Cross-attention - attn_output = self.attn2(norm_hidden_states, - context=encoder_hidden_states, - context_lens=None) - norm_hidden_states, hidden_states = self.cross_attn_residual_norm( - hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) + attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None) + norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa, + c_scale_msa) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) @@ -474,14 +455,17 @@ class WanTransformerBlock_VSA(nn.Module): self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q") self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k") self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v") - self.to_gate_compress = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_gate_compress") + self.to_gate_compress = ReplicatedLinear(dim, + dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_gate_compress") self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out") - self.attn1 = DistributedAttention_VSA( - num_heads=num_heads, - head_size=dim // num_heads, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.attn1") + self.attn1 = DistributedAttention_VSA(num_heads=num_heads, + head_size=dim // num_heads, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn1") self.hidden_dim = dim self.num_attention_heads = num_heads dim_head = dim // num_heads @@ -496,13 +480,12 @@ class WanTransformerBlock_VSA(nn.Module): print("QK Norm type not supported") raise Exception assert cross_attn_norm is True - self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( - dim, - norm_type="layer", - eps=eps, - elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32) + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32, + compute_dtype=torch.float32) # 2. Cross-attention if added_kv_proj_dim is not None: @@ -521,13 +504,12 @@ class WanTransformerBlock_VSA(nn.Module): eps=eps, quant_config=quant_config, prefix=f"{prefix}.attn2") - self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( - dim, - norm_type="layer", - eps=eps, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn") @@ -549,13 +531,11 @@ class WanTransformerBlock_VSA(nn.Module): orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 e = self.scale_shift_table + temb.float() - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( - 6, dim=1) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1) assert shift_msa.dtype == torch.float32 # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()) * - (1 + scale_msa) + shift_msa).to(orig_dtype) + norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -569,8 +549,7 @@ class WanTransformerBlock_VSA(nn.Module): query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - gate_compress = gate_compress.squeeze(1).unflatten( - 2, (self.num_attention_heads, -1)) + gate_compress = gate_compress.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) attn_output, _ = self.attn1( query, @@ -585,19 +564,15 @@ class WanTransformerBlock_VSA(nn.Module): attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.tensor([0], device=hidden_states.device) - norm_hidden_states, hidden_states = self.self_attn_residual_norm( - hidden_states, attn_output, gate_msa, null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, + null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) # 2. Cross-attention - attn_output = self.attn2(norm_hidden_states, - context=encoder_hidden_states, - context_lens=None) - norm_hidden_states, hidden_states = self.cross_attn_residual_norm( - hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) + attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None) + norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa, + c_scale_msa) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) @@ -610,14 +585,12 @@ class WanTransformerBlock_VSA(nn.Module): class WanTransformer3DModel(BaseDiT): _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions _compile_conditions = WanVideoConfig()._compile_conditions - _supported_attention_backends = WanVideoConfig( - )._supported_attention_backends + _supported_attention_backends = WanVideoConfig()._supported_attention_backends param_names_mapping = WanVideoConfig().param_names_mapping reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping - def __init__(self, config: WanVideoConfig, hf_config: dict[str, - Any]) -> None: + def __init__(self, config: WanVideoConfig, hf_config: dict[str, Any]) -> None: super().__init__(config=config, hf_config=hf_config) self.quant_config = config.quant_config @@ -630,7 +603,8 @@ class WanTransformer3DModel(BaseDiT): self.patch_size = config.patch_size self.text_len = config.text_len - assert config.num_attention_heads % get_sp_world_size() == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})" + assert config.num_attention_heads % get_sp_world_size( + ) == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})" # 1. Patch & position embedding self.patch_embedding = PatchEmbed(in_chans=config.in_channels, @@ -663,8 +637,7 @@ class WanTransformer3DModel(BaseDiT): config.added_kv_proj_dim, self._supported_attention_backends, quant_config=config.quant_config, - prefix=f"{config.prefix}.blocks.{i}") - for i in range(config.num_layers) + prefix=f"{config.prefix}.blocks.{i}") for i in range(config.num_layers) ]) # 4. Output norm & projection @@ -674,10 +647,8 @@ class WanTransformer3DModel(BaseDiT): elementwise_affine=False, dtype=torch.float32, compute_dtype=torch.float32) - self.proj_out = nn.Linear( - inner_dim, config.out_channels * math.prod(config.patch_size)) - self.scale_shift_table = nn.Parameter( - torch.randn(1, 2, inner_dim) / inner_dim**0.5) + self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) self.gradient_checkpointing = False self.__post_init__() @@ -707,27 +678,23 @@ class WanTransformer3DModel(BaseDiT): # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] - freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames, post_patch_height, - post_patch_width), - self.hidden_size, - self.num_attention_heads, - rope_dim_list, - dtype=torch.float32 if current_platform.is_mps() else torch.float64, - rope_theta=10000) - freqs_cis = (freqs_cos.to(hidden_states.device).float(), - freqs_sin.to(hidden_states.device).float()) + freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000) + freqs_cis = (freqs_cos.to(hidden_states.device).float(), freqs_sin.to(hidden_states.device).float()) hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) # Shard with padding support - returns (sharded_tensor, original_seq_len) hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1) - + current_seq_len = hidden_states.shape[1] sp_world_size = get_sp_world_size() padded_seq_len = current_seq_len * sp_world_size - # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v) if timestep.dim() == 2: @@ -756,28 +723,25 @@ class WanTransformer3DModel(BaseDiT): if encoder_hidden_states_image is not None: if encoder_hidden_states is not None: - encoder_hidden_states = torch.concat( - [encoder_hidden_states_image, encoder_hidden_states], dim=1) + encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) else: encoder_hidden_states = encoder_hidden_states_image if current_platform.is_mps() or current_platform.is_npu(): encoder_hidden_states = encoder_hidden_states.to(orig_dtype) else: - encoder_hidden_states = encoder_hidden_states # cast to orig_dtype for MPS & NPU + encoder_hidden_states = encoder_hidden_states # cast to orig_dtype for MPS & NPU assert encoder_hidden_states.dtype == orig_dtype # 4. Transformer blocks if torch.is_grad_enabled() and self.gradient_checkpointing: for block in self.blocks: - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, original_seq_len) + hidden_states = self._gradient_checkpointing_func(block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, original_seq_len) else: for block in self.blocks: - hidden_states = block(hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, original_seq_len) + hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, original_seq_len) # 5. Output norm, projection & unpatchify if temb.dim() == 3: # batch_size, seq_len, inner_dim (wan 2.2 ti2v) @@ -787,23 +751,20 @@ class WanTransformer3DModel(BaseDiT): else: # batch_size, inner_dim shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) - hidden_states = self.norm_out(hidden_states, shift, scale) # Gather and unpad in one operation - hidden_states = sequence_model_parallel_all_gather_with_unpad( - hidden_states, original_seq_len, dim=1) + hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, - post_patch_height, - post_patch_width, p_t, p_h, p_w, - -1) + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width, + p_t, p_h, p_w, -1) hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) return output + # Entry point for model registry EntryClass = WanTransformer3DModel diff --git a/fastvideo/models/dits/zimage.py b/fastvideo/models/dits/zimage.py index 969bf8be9..219327790 100644 --- a/fastvideo/models/dits/zimage.py +++ b/fastvideo/models/dits/zimage.py @@ -61,8 +61,7 @@ class TimestepEmbedder(nn.Module): def timestep_embedding(t: torch.Tensor, dim: int, max_period: int) -> torch.Tensor: with torch.amp.autocast("cuda", enabled=False): half = dim // 2 - freqs = torch.exp( - -math.log(max_period) * torch.arange(half, dtype=torch.float32, device=t.device) / half) + freqs = torch.exp(-math.log(max_period) * torch.arange(half, dtype=torch.float32, device=t.device) / half) args = t[:, None].float() * freqs[None] embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) if dim % 2: @@ -181,8 +180,7 @@ class ZImageTransformerBlock(nn.Module): self.ffn_norm2 = RMSNorm(dim, eps=norm_eps) if modulation: - self.adaLN_modulation = nn.ModuleList( - [ReplicatedLinear(min(dim, adaln_embed_dim), 4 * dim, bias=True)]) + self.adaLN_modulation = nn.ModuleList([ReplicatedLinear(min(dim, adaln_embed_dim), 4 * dim, bias=True)]) def forward( self, @@ -193,8 +191,8 @@ class ZImageTransformerBlock(nn.Module): ) -> torch.Tensor: if self.modulation: assert adaln_input is not None - scale_msa, gate_msa, scale_mlp, gate_mlp = _linear( - self.adaLN_modulation[0], adaln_input).unsqueeze(1).chunk(4, dim=2) + scale_msa, gate_msa, scale_mlp, gate_mlp = _linear(self.adaLN_modulation[0], + adaln_input).unsqueeze(1).chunk(4, dim=2) gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp @@ -258,8 +256,7 @@ class RopeEmbedder: raise ValueError("RoPE ids must have shape [sequence, number_of_axes]") if self.freqs_cis is None: self.freqs_cis = [ - freqs.to(ids.device) - for freqs in self.precompute_freqs_cis(self.axes_dims, self.axes_lens, self.theta) + freqs.to(ids.device) for freqs in self.precompute_freqs_cis(self.axes_dims, self.axes_lens, self.theta) ] elif self.freqs_cis[0].device != ids.device: self.freqs_cis = [freqs.to(ids.device) for freqs in self.freqs_cis] @@ -293,12 +290,13 @@ class ZImageTransformer2DModel(BaseDiT): self.num_channels_latents = arch.in_channels self.all_x_embedder = nn.ModuleDict({ - f"{patch_size}-{f_patch_size}": ReplicatedLinear( - f_patch_size * patch_size * patch_size * arch.in_channels, arch.dim, bias=True) + f"{patch_size}-{f_patch_size}": + ReplicatedLinear(f_patch_size * patch_size * patch_size * arch.in_channels, arch.dim, bias=True) for patch_size, f_patch_size in zip(self.all_patch_size, self.all_f_patch_size) }) self.all_final_layer = nn.ModuleDict({ - f"{patch_size}-{f_patch_size}": FinalLayer( + f"{patch_size}-{f_patch_size}": + FinalLayer( arch.dim, patch_size * patch_size * f_patch_size * self.out_channels, arch.adaln_embed_dim, @@ -334,9 +332,8 @@ class ZImageTransformer2DModel(BaseDiT): ]) self.x_pad_token = nn.Parameter(torch.empty((1, arch.dim))) self.cap_pad_token = nn.Parameter(torch.empty((1, arch.dim))) - self.layers = nn.ModuleList([ - ZImageTransformerBlock(layer_id, modulation=True, **block_kwargs) for layer_id in range(arch.n_layers) - ]) + self.layers = nn.ModuleList( + [ZImageTransformerBlock(layer_id, modulation=True, **block_kwargs) for layer_id in range(arch.n_layers)]) self.axes_dims = tuple(arch.axes_dims) self.axes_lens = tuple(arch.axes_lens) self.rope_embedder = RopeEmbedder(arch.rope_theta, self.axes_dims, self.axes_lens) @@ -386,13 +383,13 @@ class ZImageTransformer2DModel(BaseDiT): patch_size: int, f_patch_size: int, ) -> tuple[ - list[torch.Tensor], - list[torch.Tensor], - list[tuple[int, int, int]], - list[torch.Tensor], - list[torch.Tensor], - list[torch.Tensor], - list[torch.Tensor], + list[torch.Tensor], + list[torch.Tensor], + list[tuple[int, int, int]], + list[torch.Tensor], + list[torch.Tensor], + list[torch.Tensor], + list[torch.Tensor], ]: patch_height = patch_width = patch_size patch_frames = f_patch_size @@ -452,8 +449,8 @@ class ZImageTransformer2DModel(BaseDiT): device=device, ).flatten(0, 2) if image_padding: - padding_pos_ids = self.create_coordinate_grid((1, 1, 1), device=device).flatten(0, 2).repeat( - image_padding, 1) + padding_pos_ids = self.create_coordinate_grid((1, 1, 1), + device=device).flatten(0, 2).repeat(image_padding, 1) image_pos_ids.append(torch.cat([original_pos_ids, padding_pos_ids])) else: image_pos_ids.append(original_pos_ids) @@ -462,8 +459,7 @@ class ZImageTransformer2DModel(BaseDiT): torch.zeros(image_length, dtype=torch.bool, device=device), torch.ones(image_padding, dtype=torch.bool, device=device), ]) if image_padding else torch.zeros(image_length, dtype=torch.bool, device=device)) - image_out.append( - torch.cat([image, image[-1:].repeat(image_padding, 1)]) if image_padding else image) + image_out.append(torch.cat([image, image[-1:].repeat(image_padding, 1)]) if image_padding else image) return ( image_out, @@ -495,8 +491,7 @@ class ZImageTransformer2DModel(BaseDiT): ) -> tuple[list[torch.Tensor], dict]: del encoder_hidden_states_image, guidance, kwargs if model_parallel_is_initialized() and get_sp_world_size() != 1: - raise NotImplementedError( - "Z-Image masked SDPA does not support sequence parallelism; run with sp_size=1") + raise NotImplementedError("Z-Image masked SDPA does not support sequence parallelism; run with sp_size=1") if patch_size not in self.all_patch_size or f_patch_size not in self.all_f_patch_size: raise ValueError(f"unsupported patch sizes: spatial={patch_size}, temporal={f_patch_size}") if isinstance(hidden_states, torch.Tensor): @@ -551,10 +546,8 @@ class ZImageTransformer2DModel(BaseDiT): unified = [] unified_freqs_cis = [] for i, (image_length, cap_length) in enumerate(zip(image_lengths, cap_lengths)): - unified.append( - torch.cat([hidden_states[i][:image_length], encoder_hidden_states[i][:cap_length]])) - unified_freqs_cis.append( - torch.cat([image_freqs_cis[i][:image_length], cap_freqs_cis[i][:cap_length]])) + unified.append(torch.cat([hidden_states[i][:image_length], encoder_hidden_states[i][:cap_length]])) + unified_freqs_cis.append(torch.cat([image_freqs_cis[i][:image_length], cap_freqs_cis[i][:cap_length]])) unified_lengths = [image_length + cap_length for image_length, cap_length in zip(image_lengths, cap_lengths)] unified = pad_sequence(unified, batch_first=True, padding_value=0.0) unified_freqs_cis = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0) diff --git a/fastvideo/models/encoders/base.py b/fastvideo/models/encoders/base.py index 2ba2ffff1..248586d51 100644 --- a/fastvideo/models/encoders/base.py +++ b/fastvideo/models/encoders/base.py @@ -5,19 +5,14 @@ from dataclasses import field import torch from torch import nn -from fastvideo.configs.models.encoders import (BaseEncoderOutput, - ImageEncoderConfig, - TextEncoderConfig) +from fastvideo.configs.models.encoders import (BaseEncoderOutput, ImageEncoderConfig, TextEncoderConfig) from fastvideo.platforms import AttentionBackendEnum class TextEncoder(nn.Module, ABC): _fsdp_shard_conditions: list = field(default_factory=lambda: []) - _stacked_params_mapping: list[tuple[str, str, - str]] = field(default_factory=list) - _supported_attention_backends: tuple[ - AttentionBackendEnum, - ...] = TextEncoderConfig()._supported_attention_backends + _stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = TextEncoderConfig()._supported_attention_backends def __init__(self, config: TextEncoderConfig) -> None: super().__init__() @@ -25,9 +20,7 @@ class TextEncoder(nn.Module, ABC): self._fsdp_shard_conditions = config._fsdp_shard_conditions self._stacked_params_mapping = config.arch_config.stacked_params_mapping if not self.supported_attention_backends: - raise ValueError( - f"Subclass {self.__class__.__name__} must define _supported_attention_backends" - ) + raise ValueError(f"Subclass {self.__class__.__name__} must define _supported_attention_backends") @abstractmethod def forward(self, @@ -45,21 +38,16 @@ class TextEncoder(nn.Module, ABC): class ImageEncoder(nn.Module, ABC): - _supported_attention_backends: tuple[ - AttentionBackendEnum, - ...] = ImageEncoderConfig()._supported_attention_backends + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ImageEncoderConfig()._supported_attention_backends def __init__(self, config: ImageEncoderConfig) -> None: super().__init__() self.config = config if not self.supported_attention_backends: - raise ValueError( - f"Subclass {self.__class__.__name__} must define _supported_attention_backends" - ) + raise ValueError(f"Subclass {self.__class__.__name__} must define _supported_attention_backends") @abstractmethod - def forward(self, pixel_values: torch.Tensor, - **kwargs) -> BaseEncoderOutput: + def forward(self, pixel_values: torch.Tensor, **kwargs) -> BaseEncoderOutput: pass @property diff --git a/fastvideo/models/encoders/bert.py b/fastvideo/models/encoders/bert.py index 33e51f973..207af386b 100644 --- a/fastvideo/models/encoders/bert.py +++ b/fastvideo/models/encoders/bert.py @@ -17,10 +17,8 @@ class HunyuanClip(nn.Module): super().__init__() self.max_length = max_length - self.tokenizer = BertTokenizer.from_pretrained( - os.path.join(model_dir, 'tokenizer')) - self.text_encoder = BertModel.from_pretrained( - os.path.join(model_dir, 'clip_text_encoder')) + self.tokenizer = BertTokenizer.from_pretrained(os.path.join(model_dir, 'tokenizer')) + self.text_encoder = BertModel.from_pretrained(os.path.join(model_dir, 'clip_text_encoder')) @torch.no_grad def forward(self, prompts, with_mask=True): @@ -35,7 +33,6 @@ class HunyuanClip(nn.Module): ) prompt_embeds = self.text_encoder( text_inputs.input_ids.to(self.device), - attention_mask=text_inputs.attention_mask.to(self.device) - if with_mask else None, + attention_mask=text_inputs.attention_mask.to(self.device) if with_mask else None, ) return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output diff --git a/fastvideo/models/encoders/clip.py b/fastvideo/models/encoders/clip.py index 0362b6f0c..479e74477 100644 --- a/fastvideo/models/encoders/clip.py +++ b/fastvideo/models/encoders/clip.py @@ -10,12 +10,10 @@ import torch.nn as nn # from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask from fastvideo.attention import LocalAttention -from fastvideo.configs.models.encoders import (BaseEncoderOutput, - CLIPTextConfig, CLIPVisionConfig) +from fastvideo.configs.models.encoders import (BaseEncoderOutput, CLIPTextConfig, CLIPVisionConfig) from fastvideo.distributed import divide, get_tp_world_size from fastvideo.layers.activation import get_act_fn -from fastvideo.layers.linear import (ColumnParallelLinear, QKVParallelLinear, - RowParallelLinear) +from fastvideo.layers.linear import (ColumnParallelLinear, QKVParallelLinear, RowParallelLinear) from fastvideo.layers.quantization import QuantizationConfig from fastvideo.logger import init_logger from fastvideo.models.encoders.base import ImageEncoder, TextEncoder @@ -50,17 +48,13 @@ class CLIPVisionEmbeddings(nn.Module): self.num_patches = (self.image_size // self.patch_size)**2 self.num_positions = self.num_patches + 1 - self.position_embedding = nn.Embedding(self.num_positions, - self.embed_dim) - self.register_buffer("position_ids", - torch.arange(self.num_positions).expand((1, -1)), - persistent=False) + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: batch_size = pixel_values.shape[0] target_dtype = self.patch_embedding.weight.dtype - patch_embeds = self.patch_embedding(pixel_values.to( - dtype=target_dtype)) # shape = [*, width, grid, grid] + patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid] patch_embeds = patch_embeds.flatten(2).transpose(1, 2) class_embeds = self.class_embedding.expand(batch_size, 1, -1) @@ -78,14 +72,12 @@ class CLIPTextEmbeddings(nn.Module): embed_dim = config.hidden_size self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) - self.position_embedding = nn.Embedding(config.max_position_embeddings, - embed_dim) + self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim) # position_ids (1, len position emb) is contiguous in memory and exported when serialized - self.register_buffer( - "position_ids", - torch.arange(config.max_position_embeddings).expand((1, -1)), - persistent=False) + self.register_buffer("position_ids", + torch.arange(config.max_position_embeddings).expand((1, -1)), + persistent=False) def forward( self, @@ -98,16 +90,13 @@ class CLIPTextEmbeddings(nn.Module): elif inputs_embeds is not None: seq_length = inputs_embeds.shape[-2] else: - raise ValueError( - "Either input_ids or inputs_embeds must be provided.") + raise ValueError("Either input_ids or inputs_embeds must be provided.") max_position_embedding = self.position_embedding.weight.shape[0] if seq_length > max_position_embedding: - raise ValueError( - f"Sequence length must be less than max_position_embeddings (got `sequence length`: " - f"{seq_length} and max_position_embeddings: {max_position_embedding}" - ) + raise ValueError(f"Sequence length must be less than max_position_embeddings (got `sequence length`: " + f"{seq_length} and max_position_embeddings: {max_position_embedding}") if position_ids is None: position_ids = self.position_ids[:, :seq_length] @@ -136,10 +125,9 @@ class CLIPAttention(nn.Module): self.num_heads = config.num_attention_heads self.head_dim = self.embed_dim // self.num_heads if self.head_dim * self.num_heads != self.embed_dim: - raise ValueError( - "embed_dim must be divisible by num_heads " - f"(got `embed_dim`: {self.embed_dim} and `num_heads`:" - f" {self.num_heads}).") + raise ValueError("embed_dim must be divisible by num_heads " + f"(got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads}).") self.scale = self.head_dim**-0.5 if config.enable_scale else None self.dropout = config.attention_dropout @@ -161,17 +149,15 @@ class CLIPAttention(nn.Module): self.tp_size = get_tp_world_size() self.num_heads_per_partition = divide(self.num_heads, self.tp_size) - self.attn = LocalAttention( - self.num_heads_per_partition, - self.head_dim, - self.num_heads_per_partition, - softmax_scale=self.scale, - causal=config.is_causal, - supported_attention_backends=config._supported_attention_backends) + self.attn = LocalAttention(self.num_heads_per_partition, + self.head_dim, + self.num_heads_per_partition, + softmax_scale=self.scale, + causal=config.is_causal, + supported_attention_backends=config._supported_attention_backends) def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): - return tensor.view(bsz, seq_len, self.num_heads, - self.head_dim).transpose(1, 2).contiguous() + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() def forward( self, @@ -182,22 +168,15 @@ class CLIPAttention(nn.Module): qkv_states, _ = self.qkv_proj(hidden_states) query_states, key_states, value_states = qkv_states.chunk(3, dim=-1) # use flash_attn_func - query_states = query_states.reshape(query_states.shape[0], - query_states.shape[1], - self.num_heads_per_partition, + query_states = query_states.reshape(query_states.shape[0], query_states.shape[1], self.num_heads_per_partition, self.head_dim) - key_states = key_states.reshape(key_states.shape[0], - key_states.shape[1], - self.num_heads_per_partition, + key_states = key_states.reshape(key_states.shape[0], key_states.shape[1], self.num_heads_per_partition, self.head_dim) - value_states = value_states.reshape(value_states.shape[0], - value_states.shape[1], - self.num_heads_per_partition, + value_states = value_states.reshape(value_states.shape[0], value_states.shape[1], self.num_heads_per_partition, self.head_dim) attn_output = self.attn(query_states, key_states, value_states) - attn_output = attn_output.reshape( - attn_output.shape[0], attn_output.shape[1], - self.num_heads_per_partition * self.head_dim) + attn_output = attn_output.reshape(attn_output.shape[0], attn_output.shape[1], + self.num_heads_per_partition * self.head_dim) attn_output, _ = self.out_proj(attn_output) return attn_output, None @@ -247,13 +226,9 @@ class CLIPEncoderLayer(nn.Module): quant_config=quant_config, prefix=f"{prefix}.self_attn", ) - self.layer_norm1 = nn.LayerNorm(config.hidden_size, - eps=config.layer_norm_eps) - self.mlp = CLIPMLP(config, - quant_config=quant_config, - prefix=f"{prefix}.mlp") - self.layer_norm2 = nn.LayerNorm(config.hidden_size, - eps=config.layer_norm_eps) + self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.mlp = CLIPMLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp") + self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -296,15 +271,11 @@ class CLIPEncoder(nn.Module): else: num_hidden_layers = num_hidden_layers_override self.layers = nn.ModuleList([ - CLIPEncoderLayer(config=config, - quant_config=quant_config, - prefix=f"{prefix}.layers.{layer_idx}") + CLIPEncoderLayer(config=config, quant_config=quant_config, prefix=f"{prefix}.layers.{layer_idx}") for layer_idx in range(num_hidden_layers) ]) - def forward( - self, inputs_embeds: torch.Tensor, return_all_hidden_states: bool - ) -> torch.Tensor | list[torch.Tensor]: + def forward(self, inputs_embeds: torch.Tensor, return_all_hidden_states: bool) -> torch.Tensor | list[torch.Tensor]: hidden_states_pool = [inputs_embeds] hidden_states = inputs_embeds @@ -332,14 +303,12 @@ class CLIPTextTransformer(nn.Module): self.embeddings = CLIPTextEmbeddings(config) - self.encoder = CLIPEncoder( - config, - quant_config=quant_config, - num_hidden_layers_override=num_hidden_layers_override, - prefix=prefix) + self.encoder = CLIPEncoder(config, + quant_config=quant_config, + num_hidden_layers_override=num_hidden_layers_override, + prefix=prefix) - self.final_layer_norm = nn.LayerNorm(embed_dim, - eps=config.layer_norm_eps) + self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) # For `pooled_output` computation self.eos_token_id = config.eos_token_id @@ -357,8 +326,7 @@ class CLIPTextTransformer(nn.Module): """ output_hidden_states = (output_hidden_states - if output_hidden_states is not None else - self.config.output_hidden_states) + if output_hidden_states is not None else self.config.output_hidden_states) if input_ids is None: raise ValueError("You have to specify input_ids") @@ -366,8 +334,7 @@ class CLIPTextTransformer(nn.Module): input_shape = input_ids.size() input_ids = input_ids.view(-1, input_shape[-1]) - hidden_states = self.embeddings(input_ids=input_ids, - position_ids=position_ids) + hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) # CLIP's text model uses causal mask, prepare it here. # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 @@ -401,20 +368,17 @@ class CLIPTextTransformer(nn.Module): # take features from the eot embedding (eot_token is the highest number in each sequence) # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14 pooled_output = last_hidden_state[ - torch.arange(last_hidden_state.shape[0], - device=last_hidden_state.device), - input_ids.to(dtype=torch.int, device=last_hidden_state.device). - argmax(dim=-1), + torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), + input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1), ] else: # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible) pooled_output = last_hidden_state[ - torch.arange(last_hidden_state.shape[0], - device=last_hidden_state.device), + torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`) # Note: we assume each sequence (along batch dim.) contains an `eos_token_id` (e.g. prepared by the tokenizer) - (input_ids.to(dtype=torch.int, device=last_hidden_state.device - ) == self.eos_token_id).int().argmax(dim=-1), + (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id).int().argmax( + dim=-1), ] return BaseEncoderOutput( @@ -432,9 +396,7 @@ class CLIPTextModel(TextEncoder): config: CLIPTextConfig, ) -> None: super().__init__(config) - self.text_model = CLIPTextTransformer(config=config, - quant_config=config.quant_config, - prefix=config.prefix) + self.text_model = CLIPTextTransformer(config=config, quant_config=config.quant_config, prefix=config.prefix) def forward( self, @@ -454,8 +416,7 @@ class CLIPTextModel(TextEncoder): ) return outputs - def load_weights(self, weights: Iterable[tuple[str, - torch.Tensor]]) -> set[str]: + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Define mapping for stacked parameters stacked_params_mapping = [ @@ -483,8 +444,7 @@ class CLIPTextModel(TextEncoder): # Use default weight loader for all other parameters if name in params_dict: param = params_dict[name] - weight_loader = getattr(param, "weight_loader", - default_weight_loader) + weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) @@ -495,9 +455,7 @@ class CLIPTextModelWithProjection(CLIPTextModel): def __init__(self, config: CLIPTextConfig) -> None: super().__init__(config) - self.text_projection = nn.Linear( - config.hidden_size, config.projection_dim, bias=False - ) + self.text_projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False) def forward( self, @@ -558,17 +516,15 @@ class CLIPVisionTransformer(nn.Module): num_hidden_layers = config.num_hidden_layers if len(self.encoder.layers) > config.num_hidden_layers: - raise ValueError( - f"The original encoder only has {num_hidden_layers} " - f"layers, but you requested {len(self.encoder.layers)} layers.") + raise ValueError(f"The original encoder only has {num_hidden_layers} " + f"layers, but you requested {len(self.encoder.layers)} layers.") # If possible, skip post_layernorm to conserve memory if require_post_norm is None: require_post_norm = len(self.encoder.layers) == num_hidden_layers if require_post_norm: - self.post_layernorm = nn.LayerNorm(embed_dim, - eps=config.layer_norm_eps) + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) else: self.post_layernorm = None @@ -585,17 +541,14 @@ class CLIPVisionTransformer(nn.Module): # Produces either the last layer output or all of the hidden states, # depending on if we have feature_sample_layers or not - encoder_outputs = self.encoder( - inputs_embeds=hidden_states, - return_all_hidden_states=return_all_hidden_states) + encoder_outputs = self.encoder(inputs_embeds=hidden_states, return_all_hidden_states=return_all_hidden_states) if not return_all_hidden_states: encoder_outputs = encoder_outputs[0] # Handle post-norm (if applicable) and stacks feature layers if needed - encoder_outputs = resolve_visual_encoder_outputs( - encoder_outputs, feature_sample_layers, self.post_layernorm, - self.config.num_hidden_layers) + encoder_outputs = resolve_visual_encoder_outputs(encoder_outputs, feature_sample_layers, self.post_layernorm, + self.config.num_hidden_layers) return encoder_outputs @@ -607,12 +560,11 @@ class CLIPVisionModel(ImageEncoder): def __init__(self, config: CLIPVisionConfig) -> None: super().__init__(config) - self.vision_model = CLIPVisionTransformer( - config=config, - quant_config=config.quant_config, - num_hidden_layers_override=config.num_hidden_layers_override, - require_post_norm=config.require_post_norm, - prefix=f"{config.prefix}.vision_model") + self.vision_model = CLIPVisionTransformer(config=config, + quant_config=config.quant_config, + num_hidden_layers_override=config.num_hidden_layers_override, + require_post_norm=config.require_post_norm, + prefix=f"{config.prefix}.vision_model") def forward( self, @@ -620,8 +572,7 @@ class CLIPVisionModel(ImageEncoder): feature_sample_layers: list[int] | None = None, **kwargs, ) -> BaseEncoderOutput: - last_hidden_state = self.vision_model(pixel_values, - feature_sample_layers) + last_hidden_state = self.vision_model(pixel_values, feature_sample_layers) return BaseEncoderOutput(last_hidden_state=last_hidden_state) @property @@ -630,8 +581,7 @@ class CLIPVisionModel(ImageEncoder): # (TODO) Add prefix argument for filtering out weights to be loaded # ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986 - def load_weights(self, weights: Iterable[tuple[str, - torch.Tensor]]) -> set[str]: + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() @@ -641,8 +591,7 @@ class CLIPVisionModel(ImageEncoder): if name.startswith("visual_projection"): continue # post_layernorm is not needed in CLIPVisionModel - if (name.startswith("vision_model.post_layernorm") - and self.vision_model.post_layernorm is None): + if (name.startswith("vision_model.post_layernorm") and self.vision_model.post_layernorm is None): continue # omit layers when num_hidden_layers_override is set @@ -651,8 +600,7 @@ class CLIPVisionModel(ImageEncoder): if layer_idx >= layer_count: continue - for (param_name, weight_name, - shard_id) in self.config.arch_config.stacked_params_mapping: + for (param_name, weight_name, shard_id) in self.config.arch_config.stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) @@ -663,11 +611,11 @@ class CLIPVisionModel(ImageEncoder): break else: param = params_dict[name] - weight_loader = getattr(param, "weight_loader", - default_weight_loader) + weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params + # Entry point for model registry EntryClass = [CLIPTextModel, CLIPVisionModel] diff --git a/fastvideo/models/encoders/gemma.py b/fastvideo/models/encoders/gemma.py index 48f30f61e..944db35ff 100644 --- a/fastvideo/models/encoders/gemma.py +++ b/fastvideo/models/encoders/gemma.py @@ -125,9 +125,7 @@ class _BasicTransformerBlock1D(nn.Module): attention_mask: torch.Tensor | None = None, pe: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: - norm_hidden_states = torch.nn.functional.rms_norm( - hidden_states, (hidden_states.shape[-1],), eps=self.norm_eps - ) + norm_hidden_states = torch.nn.functional.rms_norm(hidden_states, (hidden_states.shape[-1], ), eps=self.norm_eps) if norm_hidden_states.ndim == 4: norm_hidden_states = norm_hidden_states.squeeze(1) @@ -140,9 +138,7 @@ class _BasicTransformerBlock1D(nn.Module): if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) - norm_hidden_states = torch.nn.functional.rms_norm( - hidden_states, (hidden_states.shape[-1],), eps=self.norm_eps - ) + norm_hidden_states = torch.nn.functional.rms_norm(hidden_states, (hidden_states.shape[-1], ), eps=self.norm_eps) ff_output = self.ff(norm_hidden_states) hidden_states = ff_output + hidden_states if hidden_states.ndim == 4: @@ -245,18 +241,15 @@ class Embeddings1DConnector(nn.Module): self.positional_embedding_max_pos = config.positional_embedding_max_pos self.rope_type = config.rope_type self.double_precision_rope = config.double_precision_rope - self.transformer_1d_blocks = nn.ModuleList( - [ - _BasicTransformerBlock1D( - dim=self.inner_dim, - heads=config.num_attention_heads, - dim_head=config.attention_head_dim, - rope_type=config.rope_type, - apply_gated_attention=config.apply_gated_attention, - ) - for _ in range(config.num_layers) - ] - ) + self.transformer_1d_blocks = nn.ModuleList([ + _BasicTransformerBlock1D( + dim=self.inner_dim, + heads=config.num_attention_heads, + dim_head=config.attention_head_dim, + rope_type=config.rope_type, + apply_gated_attention=config.apply_gated_attention, + ) for _ in range(config.num_layers) + ]) self.num_learnable_registers = config.num_learnable_registers if self.num_learnable_registers: self.learnable_registers = nn.Parameter( @@ -264,10 +257,7 @@ class Embeddings1DConnector(nn.Module): self.num_learnable_registers, self.inner_dim, dtype=torch.bfloat16, - ) - * 2.0 - - 1.0 - ) + ) * 2.0 - 1.0) def _replace_padded_with_learnable_registers( self, @@ -276,15 +266,10 @@ class Embeddings1DConnector(nn.Module): ) -> tuple[torch.Tensor, torch.Tensor]: assert hidden_states.shape[1] % self.num_learnable_registers == 0, ( f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by " - f"num_learnable_registers {self.num_learnable_registers}." - ) + f"num_learnable_registers {self.num_learnable_registers}.") - num_registers_duplications = ( - hidden_states.shape[1] // self.num_learnable_registers - ) - learnable_registers = torch.tile( - self.learnable_registers, (num_registers_duplications, 1) - ) + num_registers_duplications = (hidden_states.shape[1] // self.num_learnable_registers) + learnable_registers = torch.tile(self.learnable_registers, (num_registers_duplications, 1)) seq_len = hidden_states.shape[1] valid = attention_mask.squeeze(1).squeeze(1) >= -9000.0 # (batch, seq) @@ -318,11 +303,8 @@ class Embeddings1DConnector(nn.Module): attention_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if self.num_learnable_registers: - hidden_states, attention_mask = ( - self._replace_padded_with_learnable_registers( - hidden_states, attention_mask - ) - ) + hidden_states, attention_mask = (self._replace_padded_with_learnable_registers( + hidden_states, attention_mask)) indices_grid = torch.arange( hidden_states.shape[1], @@ -330,11 +312,8 @@ class Embeddings1DConnector(nn.Module): device=hidden_states.device, ) indices_grid = indices_grid[None, None, :] - freq_grid_generator = ( - generate_ltx_freq_grid_float64 - if self.double_precision_rope - else generate_ltx_freq_grid_pytorch - ) + freq_grid_generator = (generate_ltx_freq_grid_float64 + if self.double_precision_rope else generate_ltx_freq_grid_pytorch) freqs_cis = precompute_ltx_freqs_cis( indices_grid=indices_grid, dim=self.inner_dim, @@ -347,13 +326,9 @@ class Embeddings1DConnector(nn.Module): ) for block in self.transformer_1d_blocks: - hidden_states = block( - hidden_states, attention_mask=attention_mask, pe=freqs_cis - ) + hidden_states = block(hidden_states, attention_mask=attention_mask, pe=freqs_cis) - hidden_states = torch.nn.functional.rms_norm( - hidden_states, (hidden_states.shape[-1],), eps=1e-6 - ) + hidden_states = torch.nn.functional.rms_norm(hidden_states, (hidden_states.shape[-1], ), eps=1e-6) return hidden_states, attention_mask @@ -371,15 +346,11 @@ class LTX2GemmaTextEncoderModel(TextEncoder): # LTX-2.3 routes the caption projection before the connector and uses # separate per-modality feature extractor linears. Default (False) # keeps the single LTX-2.0 GemmaFeaturesExtractorProjLinear. - self.use_v2_feature_extractor = bool( - getattr(arch, "caption_proj_before_connector", False)) + self.use_v2_feature_extractor = bool(getattr(arch, "caption_proj_before_connector", False)) if self.use_v2_feature_extractor: - video_out_features = ( - getattr(arch, "video_feature_extractor_out_features", None) - or arch.feature_extractor_out_features) - audio_out_features = ( - getattr(arch, "audio_feature_extractor_out_features", None) - or video_out_features) + video_out_features = (getattr(arch, "video_feature_extractor_out_features", None) + or arch.feature_extractor_out_features) + audio_out_features = (getattr(arch, "audio_feature_extractor_out_features", None) or video_out_features) self.video_feature_extractor_linear = nn.Linear( arch.feature_extractor_in_features, video_out_features, @@ -396,8 +367,7 @@ class LTX2GemmaTextEncoderModel(TextEncoder): out_features=arch.feature_extractor_out_features, ) - connector_apply_gated_attention = bool( - getattr(arch, "connector_apply_gated_attention", False)) + connector_apply_gated_attention = bool(getattr(arch, "connector_apply_gated_attention", False)) video_connector_config = GemmaConnectorConfig( num_attention_heads=arch.connector_num_attention_heads, attention_head_dim=arch.connector_attention_head_dim, @@ -410,14 +380,11 @@ class LTX2GemmaTextEncoderModel(TextEncoder): apply_gated_attention=connector_apply_gated_attention, ) audio_connector_config = GemmaConnectorConfig( - num_attention_heads=getattr( - arch, "audio_connector_num_attention_heads", None) + num_attention_heads=getattr(arch, "audio_connector_num_attention_heads", None) or arch.connector_num_attention_heads, - attention_head_dim=getattr( - arch, "audio_connector_attention_head_dim", None) + attention_head_dim=getattr(arch, "audio_connector_attention_head_dim", None) or arch.connector_attention_head_dim, - num_layers=getattr(arch, "audio_connector_num_layers", None) - or arch.connector_num_layers, + num_layers=getattr(arch, "audio_connector_num_layers", None) or arch.connector_num_layers, positional_embedding_theta=arch.connector_positional_embedding_theta, positional_embedding_max_pos=arch.connector_positional_embedding_max_pos, rope_type=LTXRopeType(arch.connector_rope_type), @@ -434,9 +401,7 @@ class LTX2GemmaTextEncoderModel(TextEncoder): self._gemma_model: Gemma3ForConditionalGeneration | None = None def named_parameters(self, prefix: str = "", recurse: bool = True): - for name, param in super().named_parameters( - prefix=prefix, recurse=recurse - ): + for name, param in super().named_parameters(prefix=prefix, recurse=recurse): if name.startswith("gemma_model."): continue yield name, param @@ -475,9 +440,7 @@ class LTX2GemmaTextEncoderModel(TextEncoder): if self._gemma_model is None: gemma_path = self.gemma_model_path if not gemma_path: - raise ValueError( - "gemma_model_path must be set (expected text_encoder/gemma)." - ) + raise ValueError("gemma_model_path must be set (expected text_encoder/gemma).") dtype = getattr(torch, self.gemma_dtype, torch.bfloat16) self._gemma_model = Gemma3ForConditionalGeneration.from_pretrained( gemma_path, @@ -494,8 +457,7 @@ class LTX2GemmaTextEncoderModel(TextEncoder): if hasattr(self._gemma_model.config, "_attn_implementation"): self._gemma_model.config._attn_implementation = "sdpa" if self.use_v2_feature_extractor: - device = next( - self.video_feature_extractor_linear.parameters()).device + device = next(self.video_feature_extractor_linear.parameters()).device else: device = next(self.feature_extractor_linear.parameters()).device self._gemma_model.to(device=device) @@ -511,20 +473,15 @@ class LTX2GemmaTextEncoderModel(TextEncoder): encoded_text_features = torch.stack(hidden_states, dim=-1) if os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", ""): for idx, layer in enumerate(hidden_states): - _debug_gemma_log_line( - f"fastvideo:gemma_hidden_state_{idx}" - f":sum={layer.float().sum().item():.6f}" - ) - _debug_gemma_log_line( - "fastvideo:gemma_hidden_states_stack" - f":sum={encoded_text_features.float().sum().item():.6f}" - ) + _debug_gemma_log_line(f"fastvideo:gemma_hidden_state_{idx}" + f":sum={layer.float().sum().item():.6f}") + _debug_gemma_log_line("fastvideo:gemma_hidden_states_stack" + f":sum={encoded_text_features.float().sum().item():.6f}") encoded_text_features_dtype = encoded_text_features.dtype if self.use_v2_feature_extractor: # LTX-2.3: per-token RMS norm then separate video/audio linears. - normed_text_features = _norm_and_concat_per_token_rms( - encoded_text_features, attention_mask).to( - encoded_text_features_dtype) + normed_text_features = _norm_and_concat_per_token_rms(encoded_text_features, + attention_mask).to(encoded_text_features_dtype) video_out_dim = self.video_feature_extractor_linear.out_features audio_out_dim = self.audio_feature_extractor_linear.out_features video_features = self.video_feature_extractor_linear( @@ -544,20 +501,15 @@ class LTX2GemmaTextEncoderModel(TextEncoder): # LTX-2.0: shared min-max norm + single aggregate linear. The two # returned tensors are the same object (video == audio source). sequence_lengths = attention_mask.sum(dim=-1) - normed_text_features = _norm_and_concat_padded_batch( - encoded_text_features, sequence_lengths, padding_side=padding_side - ) - shared_features = self.feature_extractor_linear( - normed_text_features.to(encoded_text_features_dtype) - ) + normed_text_features = _norm_and_concat_padded_batch(encoded_text_features, + sequence_lengths, + padding_side=padding_side) + shared_features = self.feature_extractor_linear(normed_text_features.to(encoded_text_features_dtype)) return shared_features, shared_features - def _convert_to_additive_mask( - self, attention_mask: torch.Tensor, dtype: torch.dtype - ) -> torch.Tensor: + def _convert_to_additive_mask(self, attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: return (attention_mask - 1).to(dtype).reshape( - (attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) - ) * torch.finfo(dtype).max + (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])) * torch.finfo(dtype).max def _run_connectors( self, @@ -565,29 +517,18 @@ class LTX2GemmaTextEncoderModel(TextEncoder): audio_features: torch.Tensor | None, attention_mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - connector_attention_mask = self._convert_to_additive_mask( - attention_mask, video_features.dtype - ) - encoded, encoded_connector_attention_mask = self.embeddings_connector( - video_features, connector_attention_mask - ) + connector_attention_mask = self._convert_to_additive_mask(attention_mask, video_features.dtype) + encoded, encoded_connector_attention_mask = self.embeddings_connector(video_features, connector_attention_mask) - attention_mask = (encoded_connector_attention_mask < 0.000001).to( - torch.int64 - ) - attention_mask = attention_mask.reshape( - [encoded.shape[0], encoded.shape[1], 1] - ) + attention_mask = (encoded_connector_attention_mask < 0.000001).to(torch.int64) + attention_mask = attention_mask.reshape([encoded.shape[0], encoded.shape[1], 1]) encoded = encoded * attention_mask # For LTX-2.0 audio_features is the same tensor as video_features, so # this reproduces the prior behavior (audio connector fed the shared # features). For LTX-2.3 it is fed the separate audio features. - audio_features = (audio_features - if audio_features is not None else video_features) - encoded_for_audio, _ = self.audio_embeddings_connector( - audio_features, connector_attention_mask - ) + audio_features = (audio_features if audio_features is not None else video_features) + encoded_for_audio, _ = self.audio_embeddings_connector(audio_features, connector_attention_mask) return encoded, encoded_for_audio, attention_mask.squeeze(-1) @@ -681,25 +622,18 @@ class LTX2GemmaTextEncoderModel(TextEncoder): padding_side=self.padding_side, ) if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": - _debug_log_line( - "fastvideo:gemma_feature" - f":sum={encoded_video.float().sum().item():.6f} " - f"shape={tuple(encoded_video.shape)}" - ) - video_encoding, audio_encoding, attention_mask = self._run_connectors( - encoded_video, encoded_audio, attention_mask - ) + _debug_log_line("fastvideo:gemma_feature" + f":sum={encoded_video.float().sum().item():.6f} " + f"shape={tuple(encoded_video.shape)}") + video_encoding, audio_encoding, attention_mask = self._run_connectors(encoded_video, encoded_audio, + attention_mask) if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": - _debug_log_line( - "fastvideo:gemma_video_encoding" - f":sum={video_encoding.float().sum().item():.6f} " - f"shape={tuple(video_encoding.shape)}" - ) - _debug_log_line( - "fastvideo:gemma_audio_encoding" - f":sum={audio_encoding.float().sum().item():.6f} " - f"shape={tuple(audio_encoding.shape)}" - ) + _debug_log_line("fastvideo:gemma_video_encoding" + f":sum={video_encoding.float().sum().item():.6f} " + f"shape={tuple(video_encoding.shape)}") + _debug_log_line("fastvideo:gemma_audio_encoding" + f":sum={audio_encoding.float().sum().item():.6f} " + f"shape={tuple(audio_encoding.shape)}") hidden_states = (audio_encoding, ) if output_hidden_states else None return BaseEncoderOutput( @@ -708,9 +642,7 @@ class LTX2GemmaTextEncoderModel(TextEncoder): attention_mask=attention_mask, ) - def load_weights( - self, weights: Iterable[tuple[str, torch.Tensor]] - ) -> set[str]: + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: @@ -737,11 +669,9 @@ class LTX2GemmaTextEncoderModel(TextEncoder): elif name == "feature_extractor.aggregate_embed.weight": name = "feature_extractor_linear.aggregate_embed.weight" elif name.startswith("video_connector."): - name = name.replace( - "video_connector.", "embeddings_connector.", 1) + name = name.replace("video_connector.", "embeddings_connector.", 1) elif name.startswith("audio_connector."): - name = name.replace( - "audio_connector.", "audio_embeddings_connector.", 1) + name = name.replace("audio_connector.", "audio_embeddings_connector.", 1) if name not in params_dict: continue param = params_dict[name] @@ -766,9 +696,7 @@ def _norm_and_concat_padded_batch( start_indices = t - sequence_lengths[:, None] mask = token_indices >= start_indices else: - raise ValueError( - f"padding_side must be 'left' or 'right', got {padding_side}" - ) + raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}") mask = mask.reshape(b, t, 1, 1) eps = 1e-6 @@ -777,12 +705,8 @@ def _norm_and_concat_padded_batch( denom = (sequence_lengths * d).view(b, 1, 1, 1) mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps) - x_min = encoded_text.masked_fill(~mask, float("inf")).amin( - dim=(1, 2), keepdim=True - ) - x_max = encoded_text.masked_fill(~mask, float("-inf")).amax( - dim=(1, 2), keepdim=True - ) + x_min = encoded_text.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True) + x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True) range_ = x_max - x_min normed = 8 * (encoded_text - mean) / (range_ + eps) @@ -792,5 +716,6 @@ def _norm_and_concat_padded_batch( normed = normed.masked_fill(~mask_flattened, 0.0) return normed + # Entry point for model registry EntryClass = LTX2GemmaTextEncoderModel diff --git a/fastvideo/models/encoders/glm_image_ar_loader.py b/fastvideo/models/encoders/glm_image_ar_loader.py index 362df6fc8..7d6be0f49 100644 --- a/fastvideo/models/encoders/glm_image_ar_loader.py +++ b/fastvideo/models/encoders/glm_image_ar_loader.py @@ -14,12 +14,14 @@ logger = init_logger(__name__) class GlmImageARLoader(nn.Module): - def __init__(self, model_path: str, processor_path: str | None = None, - *, torch_dtype: torch.dtype = torch.bfloat16, + def __init__(self, + model_path: str, + processor_path: str | None = None, + *, + torch_dtype: torch.dtype = torch.bfloat16, trust_remote_code: bool = True) -> None: super().__init__() - from transformers import (AutoProcessor, - GlmImageForConditionalGeneration) + from transformers import (AutoProcessor, GlmImageForConditionalGeneration) logger.info("Loading GLM-Image AR encoder from %s", model_path) self._model = GlmImageForConditionalGeneration.from_pretrained( model_path, @@ -28,8 +30,7 @@ class GlmImageARLoader(nn.Module): ) if processor_path is not None: logger.info("Loading GLM-Image processor from %s", processor_path) - self.processor = AutoProcessor.from_pretrained( - processor_path, trust_remote_code=trust_remote_code) + self.processor = AutoProcessor.from_pretrained(processor_path, trust_remote_code=trust_remote_code) else: self.processor = None @@ -38,13 +39,11 @@ class GlmImageARLoader(nn.Module): return self._model.generate(*args, **kwargs) @torch.no_grad() - def get_image_features(self, pixel_values: torch.Tensor, - image_grid_thw: torch.Tensor) -> Any: + def get_image_features(self, pixel_values: torch.Tensor, image_grid_thw: torch.Tensor) -> Any: return self._model.get_image_features(pixel_values, image_grid_thw) @torch.no_grad() - def get_image_tokens(self, image_embeds: torch.Tensor, - image_grid_thw: torch.Tensor) -> torch.Tensor: + def get_image_tokens(self, image_embeds: torch.Tensor, image_grid_thw: torch.Tensor) -> torch.Tensor: return self._model.get_image_tokens(image_embeds, image_grid_thw) @property diff --git a/fastvideo/models/encoders/lingbotworld2_t5.py b/fastvideo/models/encoders/lingbotworld2_t5.py index 6fa619ca8..5c31cb410 100644 --- a/fastvideo/models/encoders/lingbotworld2_t5.py +++ b/fastvideo/models/encoders/lingbotworld2_t5.py @@ -35,8 +35,7 @@ def canonicalize(text: str, keep_punctuation_exact_string: str | None = None) -> if keep_punctuation_exact_string: text = keep_punctuation_exact_string.join( part.translate(str.maketrans("", "", string.punctuation)) - for part in text.split(keep_punctuation_exact_string) - ) + for part in text.split(keep_punctuation_exact_string)) else: text = text.translate(str.maketrans("", "", string.punctuation)) return " ".join(text.lower().split()) @@ -164,9 +163,8 @@ class T5RelativeEmbedding(nn.Module): rel_buckets = 0 rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos)) max_exact = num_buckets // 2 - rel_pos_large = max_exact + ( - torch.log(rel_pos.float() / max_exact) / math.log(self.max_dist / max_exact) * (num_buckets - max_exact) - ).long() + rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) / math.log(self.max_dist / max_exact) * + (num_buckets - max_exact)).long() rel_pos_large = torch.min(rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1)) rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large) return rel_buckets @@ -192,7 +190,10 @@ class T5SelfAttention(nn.Module): self.ffn = T5FeedForward(dim, dim_ffn, dropout) self.pos_embedding = None if shared_pos else T5RelativeEmbedding(num_buckets, num_heads, bidirectional=True) - def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None, pos_bias: torch.Tensor | None = None) -> torch.Tensor: + def forward(self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + pos_bias: torch.Tensor | None = None) -> torch.Tensor: """Run self-attention and feed-forward residual updates.""" e = pos_bias if self.pos_embedding is None else self.pos_embedding(x.size(1), x.size(1)) x = self._fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e)) @@ -218,20 +219,17 @@ class LingBotWorld2T5EncoderModel(TextEncoder): arch = config.arch_config self.token_embedding = nn.Embedding(arch.vocab_size, arch.dim) self.dropout = nn.Dropout(arch.dropout) - self.blocks = nn.ModuleList( - [ - T5SelfAttention( - arch.dim, - arch.dim_attn, - arch.dim_ffn, - arch.num_heads, - arch.num_buckets, - shared_pos=False, - dropout=arch.dropout, - ) - for _ in range(arch.num_layers) - ] - ) + self.blocks = nn.ModuleList([ + T5SelfAttention( + arch.dim, + arch.dim_attn, + arch.dim_ffn, + arch.num_heads, + arch.num_buckets, + shared_pos=False, + dropout=arch.dropout, + ) for _ in range(arch.num_layers) + ]) self.norm = T5LayerNorm(arch.dim) def forward( diff --git a/fastvideo/models/encoders/llama.py b/fastvideo/models/encoders/llama.py index 88c548aec..f04867bc5 100644 --- a/fastvideo/models/encoders/llama.py +++ b/fastvideo/models/encoders/llama.py @@ -36,14 +36,12 @@ from fastvideo.configs.models.encoders import BaseEncoderOutput, LlamaConfig from fastvideo.distributed import get_tp_world_size from fastvideo.layers.activation import SiluAndMul from fastvideo.layers.layernorm import RMSNorm -from fastvideo.layers.linear import (MergedColumnParallelLinear, - QKVParallelLinear, RowParallelLinear) +from fastvideo.layers.linear import (MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear) from fastvideo.layers.quantization import QuantizationConfig from fastvideo.layers.rotary_embedding import get_rope from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding from fastvideo.models.encoders.base import TextEncoder -from fastvideo.models.loader.weight_utils import (default_weight_loader, - maybe_remap_kv_scale_name) +from fastvideo.models.loader.weight_utils import (default_weight_loader, maybe_remap_kv_scale_name) class LlamaMLP(nn.Module): @@ -117,8 +115,7 @@ class LlamaAttention(nn.Module): assert tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) # MistralConfig has an optional head_dim introduced by Mistral-Nemo - self.head_dim = getattr(config, "head_dim", - self.hidden_size // self.total_num_heads) + self.head_dim = getattr(config, "head_dim", self.hidden_size // self.total_num_heads) # Phi models introduced a partial_rotary_factor parameter in the config partial_rotary_factor = getattr(config, "partial_rotary_factor", 1) self.rotary_dim = int(partial_rotary_factor * self.head_dim) @@ -147,8 +144,7 @@ class LlamaAttention(nn.Module): ) is_neox_style = True - is_gguf = quant_config and hasattr( - quant_config, "get_name") and quant_config.get_name() == "gguf" + is_gguf = quant_config and hasattr(quant_config, "get_name") and quant_config.get_name() == "gguf" if is_gguf and config.model_type == "llama": is_neox_style = False @@ -161,13 +157,12 @@ class LlamaAttention(nn.Module): is_neox_style=is_neox_style, ) - self.attn = LocalAttention( - self.num_heads, - self.head_dim, - self.num_kv_heads, - softmax_scale=self.scaling, - causal=True, - supported_attention_backends=config._supported_attention_backends) + self.attn = LocalAttention(self.num_heads, + self.head_dim, + self.num_kv_heads, + softmax_scale=self.scaling, + causal=True, + supported_attention_backends=config._supported_attention_backends) def forward( self, @@ -190,8 +185,7 @@ class LlamaAttention(nn.Module): # import pdb; pdb.set_trace() # attn_output = flash_attn_func(q, k, v, softmax_scale=self.scaling, causal=True) attn_output = self.attn(q, k, v) - attn_output = attn_output.reshape(batch_size, seq_len, - self.num_heads * self.head_dim) + attn_output = attn_output.reshape(batch_size, seq_len, self.num_heads * self.head_dim) output, _ = self.o_proj(attn_output) return output @@ -209,16 +203,12 @@ class LlamaDecoderLayer(nn.Module): self.hidden_size = config.hidden_size rope_theta = getattr(config, "rope_theta", 10000) rope_scaling = getattr(config, "rope_scaling", None) - if rope_scaling is not None and getattr( - config, "original_max_position_embeddings", None): - rope_scaling["original_max_position_embeddings"] = ( - config.original_max_position_embeddings) - max_position_embeddings = getattr(config, "max_position_embeddings", - 8192) + if rope_scaling is not None and getattr(config, "original_max_position_embeddings", None): + rope_scaling["original_max_position_embeddings"] = (config.original_max_position_embeddings) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias # Support internlm/internlm-7b with bias - attention_bias = getattr(config, "attention_bias", False) or getattr( - config, "bias", False) + attention_bias = getattr(config, "attention_bias", False) or getattr(config, "bias", False) bias_o_proj = attention_bias # support internlm/internlm3-8b with qkv_bias if hasattr(config, 'qkv_bias'): @@ -228,8 +218,7 @@ class LlamaDecoderLayer(nn.Module): config=config, hidden_size=self.hidden_size, num_heads=config.num_attention_heads, - num_kv_heads=getattr(config, "num_key_value_heads", - config.num_attention_heads), + num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads), rope_theta=rope_theta, rope_scaling=rope_scaling, max_position_embeddings=max_position_embeddings, @@ -246,10 +235,8 @@ class LlamaDecoderLayer(nn.Module): bias=getattr(config, "mlp_bias", False), prefix=f"{prefix}.mlp", ) - self.input_layernorm = RMSNorm(config.hidden_size, - eps=config.rms_norm_eps) - self.post_attention_layernorm = RMSNorm(config.hidden_size, - eps=config.rms_norm_eps) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, @@ -262,15 +249,12 @@ class LlamaDecoderLayer(nn.Module): residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: - hidden_states, residual = self.input_layernorm( - hidden_states, residual) + hidden_states, residual = self.input_layernorm(hidden_states, residual) - hidden_states = self.self_attn(positions=positions, - hidden_states=hidden_states) + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) # Fully Connected - hidden_states, residual = self.post_attention_layernorm( - hidden_states, residual) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) hidden_states = self.mlp(hidden_states) return hidden_states, residual @@ -306,9 +290,7 @@ class LlamaModel(TextEncoder): ) self.layers = nn.ModuleList([ - LlamaDecoderLayer(config=config, - quant_config=config.quant_config, - prefix=f"{config.prefix}.layers.{i}") + LlamaDecoderLayer(config=config, quant_config=config.quant_config, prefix=f"{config.prefix}.layers.{i}") for i in range(config.num_hidden_layers) ]) @@ -327,26 +309,19 @@ class LlamaModel(TextEncoder): **kwargs, ) -> BaseEncoderOutput: output_hidden_states = (output_hidden_states - if output_hidden_states is not None else - self.config.output_hidden_states) + if output_hidden_states is not None else self.config.output_hidden_states) hidden_states = inputs_embeds if inputs_embeds is not None else self.get_input_embeddings(input_ids) residual = None if position_ids is None: - position_ids = torch.arange( - 0, hidden_states.shape[1], - device=hidden_states.device).unsqueeze(0) + position_ids = torch.arange(0, hidden_states.shape[1], device=hidden_states.device).unsqueeze(0) - all_hidden_states: tuple[Any, ...] | None = ( - ) if output_hidden_states else None + all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None for layer in self.layers: if all_hidden_states is not None: # TODO - all_hidden_states += ( - hidden_states, ) if residual is None else (hidden_states + - residual, ) - hidden_states, residual = layer(position_ids, hidden_states, - residual) + all_hidden_states += (hidden_states, ) if residual is None else (hidden_states + residual, ) + hidden_states, residual = layer(position_ids, hidden_states, residual) hidden_states, _ = self.norm(hidden_states, residual) @@ -365,16 +340,14 @@ class LlamaModel(TextEncoder): return output - def load_weights(self, weights: Iterable[tuple[str, - torch.Tensor]]) -> set[str]: + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue - if ("rotary_emb.cos_cached" in name - or "rotary_emb.sin_cached" in name): + if ("rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name): # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue @@ -391,8 +364,7 @@ class LlamaModel(TextEncoder): # continue if "scale" in name: # Remapping the name of FP8 kv-scale. - kv_scale_name: str | None = maybe_remap_kv_scale_name( - name, params_dict) + kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict) if kv_scale_name is None: continue else: @@ -421,11 +393,11 @@ class LlamaModel(TextEncoder): continue param = params_dict[name] - weight_loader = getattr(param, "weight_loader", - default_weight_loader) + weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params + # Entry point for model registry EntryClass = LlamaModel diff --git a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py index 6c9d72691..b4570dcc2 100644 --- a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py +++ b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py @@ -52,6 +52,7 @@ class MiniMaxH3Qwen3VLTextRotaryEmbedding(nn.Module): class MiniMaxH3Qwen3VLTextAttention(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig, prefix: str) -> None: super().__init__() tp_size = get_tp_world_size() @@ -60,10 +61,8 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module): if self.total_num_heads % tp_size: raise ValueError(f"Qwen3-VL attention heads {self.total_num_heads} are not divisible by TP={tp_size}") if tp_size > self.total_num_kv_heads: - raise ValueError( - f"Qwen3-VL native K/V projections require TP={tp_size} to be no larger than " - f"the {self.total_num_kv_heads} KV heads" - ) + raise ValueError(f"Qwen3-VL native K/V projections require TP={tp_size} to be no larger than " + f"the {self.total_num_kv_heads} KV heads") if self.total_num_kv_heads % tp_size: raise ValueError(f"Qwen3-VL KV heads {self.total_num_kv_heads} are not divisible by TP={tp_size}") @@ -136,10 +135,10 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module): else: if self.num_heads != self.num_kv_heads: groups = self.num_heads // self.num_kv_heads - key = key[:, :, None].expand(-1, -1, groups, -1, -1).reshape( - batch_size, self.num_heads, sequence_length, self.head_dim) - value = value[:, :, None].expand(-1, -1, groups, -1, -1).reshape( - batch_size, self.num_heads, sequence_length, self.head_dim) + key = key[:, :, None].expand(-1, -1, groups, -1, -1).reshape(batch_size, self.num_heads, + sequence_length, self.head_dim) + value = value[:, :, None].expand(-1, -1, groups, -1, -1).reshape(batch_size, self.num_heads, + sequence_length, self.head_dim) causal = torch.ones(sequence_length, sequence_length, device=query.device, dtype=torch.bool).tril() key_mask = attention_mask.to(device=query.device, dtype=torch.bool) mask = causal[None, None] & key_mask[:, None, None] @@ -158,6 +157,7 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module): class MiniMaxH3Qwen3VLTextMLP(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig, prefix: str) -> None: super().__init__() quant_config = getattr(config, "quant_config", None) @@ -182,6 +182,7 @@ class MiniMaxH3Qwen3VLTextMLP(nn.Module): quant_config=quant_config, prefix=f"{prefix}.down_proj", ) + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: gate, _ = self.gate_proj(hidden_states) up, _ = self.up_proj(hidden_states) @@ -191,6 +192,7 @@ class MiniMaxH3Qwen3VLTextMLP(nn.Module): class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig, prefix: str) -> None: super().__init__() self.self_attn = MiniMaxH3Qwen3VLTextAttention(config, prefix=f"{prefix}.self_attn") @@ -215,6 +217,7 @@ class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module): class MiniMaxH3Qwen3VLLanguageModel(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: super().__init__() quant_config = getattr(config, "quant_config", None) @@ -226,8 +229,7 @@ class MiniMaxH3Qwen3VLLanguageModel(nn.Module): ) self.layers = nn.ModuleList( MiniMaxH3Qwen3VLTextDecoderLayer(config, prefix=f"{config.prefix}.language_model.layers.{index}") - for index in range(config.num_hidden_layers) - ) + for index in range(config.num_hidden_layers)) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = MiniMaxH3Qwen3VLTextRotaryEmbedding(config) @@ -247,7 +249,7 @@ class MiniMaxH3Qwen3VLLanguageModel(nn.Module): all_hidden_states: tuple[torch.Tensor, ...] | None = () if output_hidden_states else None for layer_index, layer in enumerate(self.layers): if all_hidden_states is not None: - all_hidden_states += (hidden_states,) + all_hidden_states += (hidden_states, ) hidden_states = layer(hidden_states, position_embeddings, attention_mask) if deepstack_visual_embeds is not None and layer_index < len(deepstack_visual_embeds): if visual_pos_masks is None: @@ -258,11 +260,12 @@ class MiniMaxH3Qwen3VLLanguageModel(nn.Module): hidden_states[mask] = updated hidden_states = self.norm(hidden_states) if all_hidden_states is not None: - all_hidden_states += (hidden_states,) + all_hidden_states += (hidden_states, ) return BaseEncoderOutput(last_hidden_state=hidden_states, hidden_states=all_hidden_states) class MiniMaxH3Qwen3VLVisionPatchEmbed(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: super().__init__() self.in_channels = config.vision_in_channels @@ -278,6 +281,7 @@ class MiniMaxH3Qwen3VLVisionPatchEmbed(nn.Module): class MiniMaxH3Qwen3VLVisionRotaryEmbedding(nn.Module): + def __init__(self, dimension: int) -> None: super().__init__() # Transformers initializes this non-persistent constant on CPU before @@ -295,6 +299,7 @@ class MiniMaxH3Qwen3VLVisionRotaryEmbedding(nn.Module): class MiniMaxH3Qwen3VLVisionPatchMerger(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig, use_postshuffle_norm: bool) -> None: super().__init__() self.hidden_size = config.vision_hidden_size * config.vision_spatial_merge_size**2 @@ -315,6 +320,7 @@ class MiniMaxH3Qwen3VLVisionPatchMerger(nn.Module): class MiniMaxH3Qwen3VLVisionMLP(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: super().__init__() self.linear_fc1 = nn.Linear(config.vision_hidden_size, config.vision_intermediate_size, bias=True) @@ -325,6 +331,7 @@ class MiniMaxH3Qwen3VLVisionMLP(nn.Module): class MiniMaxH3Qwen3VLVisionAttention(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: super().__init__() self.num_heads = config.vision_num_heads @@ -367,6 +374,7 @@ class MiniMaxH3Qwen3VLVisionAttention(nn.Module): class MiniMaxH3Qwen3VLVisionBlock(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: super().__init__() self.norm1 = nn.LayerNorm(config.vision_hidden_size, eps=1e-6) @@ -385,6 +393,7 @@ class MiniMaxH3Qwen3VLVisionBlock(nn.Module): class MiniMaxH3Qwen3VLVisionModel(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: super().__init__() self.config = config @@ -398,9 +407,7 @@ class MiniMaxH3Qwen3VLVisionModel(nn.Module): self.merger = MiniMaxH3Qwen3VLVisionPatchMerger(config, use_postshuffle_norm=False) self.deepstack_visual_indexes = tuple(config.vision_deepstack_visual_indexes) self.deepstack_merger_list = nn.ModuleList( - MiniMaxH3Qwen3VLVisionPatchMerger(config, use_postshuffle_norm=True) - for _ in self.deepstack_visual_indexes - ) + MiniMaxH3Qwen3VLVisionPatchMerger(config, use_postshuffle_norm=True) for _ in self.deepstack_visual_indexes) def _rotary_positions(self, grid_thw: torch.Tensor) -> torch.Tensor: max_height_width = int(grid_thw[:, 1:].max().item()) @@ -462,7 +469,9 @@ class MiniMaxH3Qwen3VLVisionModel(nn.Module): weight_lists[index].extend(weights[index].tolist()) index_tensor = torch.tensor(index_lists, dtype=torch.long, device=self.pos_embed.weight.device) - weight_tensor = torch.tensor(weight_lists, dtype=self.pos_embed.weight.dtype, device=self.pos_embed.weight.device) + weight_tensor = torch.tensor(weight_lists, + dtype=self.pos_embed.weight.dtype, + device=self.pos_embed.weight.device) embeddings = self.pos_embed(index_tensor) * weight_tensor[:, :, None] embeddings = (embeddings[0] + embeddings[1] + embeddings[2] + embeddings[3]).split(patch_counts) permuted = [] @@ -478,9 +487,9 @@ class MiniMaxH3Qwen3VLVisionModel(nn.Module): rotary = self._rotary_positions(grid_thw).reshape(hidden_states.shape[0], -1) embedding = torch.cat((rotary, rotary), dim=-1) position_embeddings = (embedding.cos(), embedding.sin()) - sequence_lengths = [int(value) for value in torch.repeat_interleave( - grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] - )] + sequence_lengths = [ + int(value) for value in torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]) + ] deepstack_features = [] for layer_index, block in enumerate(self.blocks): hidden_states = block(hidden_states, sequence_lengths, position_embeddings) @@ -518,8 +527,8 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder): ) -> torch.Tensor: if image_grid_thw is None and video_grid_thw is None: if attention_mask is None: - return torch.arange(input_ids.shape[1], device=input_ids.device).view(1, 1, -1).expand( - 3, input_ids.shape[0], -1) + return torch.arange(input_ids.shape[1], + device=input_ids.device).view(1, 1, -1).expand(3, input_ids.shape[0], -1) positions = attention_mask.long().cumsum(-1) - 1 positions.masked_fill_(attention_mask == 0, 1) return positions.unsqueeze(0).expand(3, -1, -1) @@ -528,7 +537,11 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder): video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0).clone() video_grid_thw[:, 0] = 1 mask = torch.ones_like(input_ids) if attention_mask is None else attention_mask - position_ids = torch.ones(3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device) + position_ids = torch.ones(3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device) image_index = 0 video_index = 0 merge = self.config.vision_spatial_merge_size @@ -593,10 +606,8 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder): mask = input_ids == token_id expanded = mask.unsqueeze(-1).expand_as(inputs_embeds) if inputs_embeds[expanded].numel() != features.numel(): - raise ValueError( - f"Qwen3-VL {label} features and placeholder tokens do not match: " - f"tokens={int(mask.sum())}, features={features.shape[0]}" - ) + raise ValueError(f"Qwen3-VL {label} features and placeholder tokens do not match: " + f"tokens={int(mask.sum())}, features={features.shape[0]}") return mask def forward( @@ -631,16 +642,16 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder): raise ValueError("pixel_values require input_ids and image_grid_thw") image_features, image_deepstack = self._visual_features(pixel_values, image_grid_thw) image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) - image_mask = self._placeholder_mask( - input_ids, inputs_embeds, self.config.image_token_id, image_features, "image") + image_mask = self._placeholder_mask(input_ids, inputs_embeds, self.config.image_token_id, image_features, + "image") inputs_embeds = inputs_embeds.masked_scatter(image_mask.unsqueeze(-1), image_features) if pixel_values_videos is not None: if input_ids is None or video_grid_thw is None: raise ValueError("pixel_values_videos require input_ids and video_grid_thw") video_features, video_deepstack = self._visual_features(pixel_values_videos, video_grid_thw) video_features = video_features.to(inputs_embeds.device, inputs_embeds.dtype) - video_mask = self._placeholder_mask( - input_ids, inputs_embeds, self.config.video_token_id, video_features, "video") + video_mask = self._placeholder_mask(input_ids, inputs_embeds, self.config.video_token_id, video_features, + "video") inputs_embeds = inputs_embeds.masked_scatter(video_mask.unsqueeze(-1), video_features) visual_mask = None @@ -666,8 +677,9 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder): if position_ids is None: if input_ids is None: sequence_length = inputs_embeds.shape[1] - position_ids = torch.arange(sequence_length, device=inputs_embeds.device).view(1, 1, -1).expand( - 3, inputs_embeds.shape[0], -1) + position_ids = torch.arange(sequence_length, + device=inputs_embeds.device).view(1, 1, + -1).expand(3, inputs_embeds.shape[0], -1) else: position_ids = self._get_rope_index(input_ids, image_grid_thw, video_grid_thw, attention_mask) output_hidden_states = self.config.output_hidden_states if output_hidden_states is None else output_hidden_states diff --git a/fastvideo/models/encoders/mistral3.py b/fastvideo/models/encoders/mistral3.py index c04e786f3..e394361c1 100644 --- a/fastvideo/models/encoders/mistral3.py +++ b/fastvideo/models/encoders/mistral3.py @@ -39,10 +39,8 @@ class Mistral3ForConditionalGeneration(TextEncoder): return model def forward(self, *args: Any, **kwargs: Any) -> Any: - raise NotImplementedError( - "Mistral3ForConditionalGeneration is loaded through Transformers " - "via from_pretrained_local()." - ) + raise NotImplementedError("Mistral3ForConditionalGeneration is loaded through Transformers " + "via from_pretrained_local().") EntryClass = Mistral3ForConditionalGeneration diff --git a/fastvideo/models/encoders/qwen2_5.py b/fastvideo/models/encoders/qwen2_5.py index 724e22889..a94840c18 100644 --- a/fastvideo/models/encoders/qwen2_5.py +++ b/fastvideo/models/encoders/qwen2_5.py @@ -22,6 +22,7 @@ from fastvideo.logger import init_logger logger = init_logger(__name__) + def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, @@ -33,6 +34,7 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + def sdpa_attention_forward( module: torch.nn.Module, query: torch.Tensor, @@ -45,17 +47,15 @@ def sdpa_attention_forward( **kwargs, ) -> tuple[torch.Tensor, None]: if kwargs.get("output_attentions", False) or kwargs.get("head_mask") is not None: - logger.warning_once( - "`sdpa` attention does not support `output_attentions=True` or `head_mask`." - " Please set your attention to `eager` if you want any of these features." - ) + logger.warning_once("`sdpa` attention does not support `output_attentions=True` or `head_mask`." + " Please set your attention to `eager` if you want any of these features.") if hasattr(module, "num_key_value_groups"): key = repeat_kv(key, module.num_key_value_groups) value = repeat_kv(value, module.num_key_value_groups) if attention_mask is not None and attention_mask.ndim == 4: - attention_mask = attention_mask[:, :, :, : key.shape[-2]] + attention_mask = attention_mask[:, :, :, :key.shape[-2]] # If attention_mask is not None, convert it to boolean type if attention_mask is not None and attention_mask.dtype != torch.bool: attention_mask = attention_mask.bool() @@ -73,21 +73,18 @@ def sdpa_attention_forward( return attn_output, None + def rotate_half(x): """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] + x1 = x[..., :x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1): mrope_section = [s * 2 for s in mrope_section] - cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze( - unsqueeze_dim - ) - sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze( - unsqueeze_dim - ) + cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim) + sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) @@ -95,23 +92,22 @@ def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim class Qwen2_5_VLRotaryEmbedding(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, device=None): super().__init__() self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config - + self.rope_type = config.rope_scaling.get("rope_type", "default") self.base = config.rope_theta - + # Simplified initialization head_dim = config.hidden_size // config.num_attention_heads dim = head_dim self.attention_scaling = 1.0 - inv_freq = 1.0 / ( - self.base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim) - ) + inv_freq = 1.0 / (self.base**(torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = inv_freq @@ -132,6 +128,7 @@ class Qwen2_5_VLRotaryEmbedding(nn.Module): class Qwen2_5_VLMLP(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, quant_config: QuantizationConfig | None = None, prefix: str = ""): super().__init__() self.gate_up_proj = MergedColumnParallelLinear( @@ -158,7 +155,12 @@ class Qwen2_5_VLMLP(nn.Module): class Qwen2_5_VLAttention(nn.Module): - def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int, quant_config: QuantizationConfig | None = None, prefix: str = ""): + + def __init__(self, + config: Qwen2_5_VLConfig, + layer_idx: int, + quant_config: QuantizationConfig | None = None, + prefix: str = ""): super().__init__() self.config = config self.layer_idx = layer_idx @@ -167,22 +169,22 @@ class Qwen2_5_VLAttention(nn.Module): self.head_dim = self.hidden_size // self.num_heads self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads - + tp_size = get_tp_world_size() self.total_num_heads = self.num_heads assert self.total_num_heads % tp_size == 0 self.num_heads = self.total_num_heads // tp_size - + self.total_num_kv_heads = self.num_key_value_heads if self.total_num_kv_heads >= tp_size: - assert self.total_num_kv_heads % tp_size == 0 - self.num_kv_heads = self.total_num_kv_heads // tp_size + assert self.total_num_kv_heads % tp_size == 0 + self.num_kv_heads = self.total_num_kv_heads // tp_size else: - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = 1 + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = 1 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - + self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 @@ -196,7 +198,7 @@ class Qwen2_5_VLAttention(nn.Module): quant_config=quant_config, prefix=f"{prefix}.qkv_proj", ) - + self.o_proj = RowParallelLinear( input_size=self.total_num_heads * self.head_dim, output_size=self.hidden_size, @@ -204,7 +206,7 @@ class Qwen2_5_VLAttention(nn.Module): quant_config=quant_config, prefix=f"{prefix}.o_proj", ) - + self.layer_type = config.layer_types[layer_idx] if config.layer_types else "full_attention" self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None @@ -226,17 +228,29 @@ class Qwen2_5_VLAttention(nn.Module): value_states = value_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2) cos, sin = position_embeddings - query_states, key_states = apply_multimodal_rotary_pos_emb( - query_states, key_states, cos, sin, self.config.rope_scaling["mrope_section"] - ) + query_states, key_states = apply_multimodal_rotary_pos_emb(query_states, key_states, cos, sin, + self.config.rope_scaling["mrope_section"]) - attn_output = sdpa_attention_forward(self, query_states, key_states, value_states, attention_mask, dropout=self.config.attention_dropout, scaling=self.scaling, is_causal=False)[0].reshape(bsz, q_len, -1) + attn_output = sdpa_attention_forward(self, + query_states, + key_states, + value_states, + attention_mask, + dropout=self.config.attention_dropout, + scaling=self.scaling, + is_causal=False)[0].reshape(bsz, q_len, -1) attn_output, _ = self.o_proj(attn_output) - + return attn_output + class Qwen2_5_VLDecoderLayer(nn.Module): - def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int, quant_config: QuantizationConfig | None = None, prefix: str = ""): + + def __init__(self, + config: Qwen2_5_VLConfig, + layer_idx: int, + quant_config: QuantizationConfig | None = None, + prefix: str = ""): super().__init__() self.self_attn = Qwen2_5_VLAttention(config, layer_idx, quant_config=quant_config, prefix=f"{prefix}.self_attn") self.mlp = Qwen2_5_VLMLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp") @@ -251,10 +265,10 @@ class Qwen2_5_VLDecoderLayer(nn.Module): position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, output_attentions: Optional[bool] = False, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: - + residual = hidden_states hidden_states = self.input_layernorm(hidden_states) - + hidden_states = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, @@ -263,33 +277,36 @@ class Qwen2_5_VLDecoderLayer(nn.Module): output_attentions=output_attentions, ) hidden_states = residual + hidden_states - + residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states - - outputs = (hidden_states,) - + + outputs = (hidden_states, ) + return outputs + class Qwen2_5_VLTextModel(TextEncoder): + def __init__(self, config: Qwen2_5_VLConfig): super().__init__(config) quant_config = None - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - org_num_embeddings=config.vocab_size - ) - + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size) + self.layers = nn.ModuleList([ - Qwen2_5_VLDecoderLayer(config, layer_idx, quant_config=quant_config, prefix=f"{config.prefix}.layers.{layer_idx}") + Qwen2_5_VLDecoderLayer(config, + layer_idx, + quant_config=quant_config, + prefix=f"{config.prefix}.layers.{layer_idx}") for layer_idx in range(config.num_hidden_layers) ]) - + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config) @@ -305,15 +322,14 @@ class Qwen2_5_VLTextModel(TextEncoder): output_hidden_states: Optional[bool] = None, **kwargs, ) -> BaseEncoderOutput: - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else self.config.output_hidden_states) + if inputs_embeds is None: inputs_embeds = self.get_input_embeddings(input_ids) - + hidden_states = inputs_embeds - + if position_ids is None: seq_length = hidden_states.shape[1] cache_position = torch.arange(seq_length, device=hidden_states.device) @@ -328,13 +344,13 @@ class Qwen2_5_VLTextModel(TextEncoder): } position_embeddings = self.rotary_emb(hidden_states, position_ids) - + all_hidden_states = () if output_hidden_states else None - + for decoder_layer in self.layers: if output_hidden_states: - all_hidden_states += (hidden_states,) - + all_hidden_states += (hidden_states, ) + layer_outputs = decoder_layer( hidden_states, attention_mask=sdpa_mask(**mask_kwargs), @@ -343,12 +359,12 @@ class Qwen2_5_VLTextModel(TextEncoder): output_attentions=False, ) hidden_states = layer_outputs[0] - + hidden_states = self.norm(hidden_states) - + if output_hidden_states: - all_hidden_states += (hidden_states,) - + all_hidden_states += (hidden_states, ) + return BaseEncoderOutput( last_hidden_state=hidden_states, hidden_states=all_hidden_states, @@ -386,5 +402,6 @@ class Qwen2_5_VLTextModel(TextEncoder): loaded_params.add(name) return loaded_params + # Entry point for model registry EntryClass = Qwen2_5_VLTextModel diff --git a/fastvideo/models/encoders/qwen2_5_vl_custom.py b/fastvideo/models/encoders/qwen2_5_vl_custom.py index fbb5c6d5d..64f3aef50 100644 --- a/fastvideo/models/encoders/qwen2_5_vl_custom.py +++ b/fastvideo/models/encoders/qwen2_5_vl_custom.py @@ -59,7 +59,6 @@ else: flash_attn_varlen_func = None apply_rotary_emb = None - if is_flash_attn_2_available(): from transformers.modeling_flash_attention_utils import _flash_attention_forward else: @@ -95,6 +94,7 @@ def _resolve_torch_dtype(dtype, default: torch.dtype = torch.float32) -> torch.d class Qwen2_5_VLMLP(nn.Module): + def __init__(self, config, bias: bool = False): super().__init__() self.hidden_size = config.hidden_size @@ -109,6 +109,7 @@ class Qwen2_5_VLMLP(nn.Module): class Qwen2_5_VisionPatchEmbed(nn.Module): + def __init__( self, patch_size: int = 14, @@ -145,15 +146,16 @@ class Qwen2_5_VisionPatchEmbed(nn.Module): class Qwen2_5_VisionRotaryEmbedding(nn.Module): + def __init__(self, dim: int, theta: float = 10000.0) -> None: super().__init__() - inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + inv_freq = 1.0 / (theta**(torch.arange(0, dim, 2, dtype=torch.float) / dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) self.dim = dim def init_weights(self, buffer_device: Optional[torch.device] = None): device = self.inv_freq.device if buffer_device is None else buffer_device - self.inv_freq = 1.0 / (10000.0 ** (torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim)).to(device) + self.inv_freq = 1.0 / (10000.0**(torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim)).to(device) def forward(self, seqlen: int) -> torch.Tensor: seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype) @@ -162,6 +164,7 @@ class Qwen2_5_VisionRotaryEmbedding(nn.Module): class Qwen2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): """ Qwen2RMSNorm is equivalent to T5LayerNorm @@ -182,6 +185,7 @@ class Qwen2RMSNorm(nn.Module): class Qwen2_5_VLPatchMerger(nn.Module): + def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None: super().__init__() self.hidden_size = context_dim * (spatial_merge_size**2) @@ -209,6 +213,7 @@ def apply_rotary_pos_emb_flashatt(tensor: torch.Tensor, freqs: torch.Tensor) -> class Qwen2_5_VLVisionFlashAttention2(nn.Module): + def __init__(self, dim: int, num_heads: int = 16) -> None: super().__init__() self.num_heads = num_heads @@ -227,17 +232,16 @@ class Qwen2_5_VLVisionFlashAttention2(nn.Module): k = apply_rotary_pos_emb_flashatt(k.unsqueeze(0), rotary_pos_emb).squeeze(0) max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() - attn_output = flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen).reshape( - seq_length, -1 - ) + attn_output = flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, + max_seqlen).reshape(seq_length, -1) attn_output = self.proj(attn_output) return attn_output def rotate_half(x): """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] + x1 = x[..., :x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) @@ -254,6 +258,7 @@ def apply_rotary_pos_emb_vision(tensor: torch.Tensor, freqs: torch.Tensor) -> to class Qwen2_5_VLVisionSdpaAttention(nn.Module): + def __init__(self, dim: int, num_heads: int = 16) -> None: super().__init__() self.num_heads = num_heads @@ -275,8 +280,8 @@ class Qwen2_5_VLVisionSdpaAttention(nn.Module): for i in range(1, len(cu_seqlens)): attention_mask[ ..., - cu_seqlens[i - 1] : cu_seqlens[i], - cu_seqlens[i - 1] : cu_seqlens[i], + cu_seqlens[i - 1]:cu_seqlens[i], + cu_seqlens[i - 1]:cu_seqlens[i], ] = True q = q.transpose(0, 1) k = k.transpose(0, 1) @@ -295,18 +300,17 @@ QWEN2_5_VL_VISION_ATTENTION_CLASSES = { # If FlashAttention2 is not available, transparently fall back to SDPA. if not is_flash_attn_2_available(): - QWEN2_5_VL_VISION_ATTENTION_CLASSES["flash_attention_2"] = ( - Qwen2_5_VLVisionSdpaAttention - ) + QWEN2_5_VL_VISION_ATTENTION_CLASSES["flash_attention_2"] = (Qwen2_5_VLVisionSdpaAttention) + class Qwen2_5_VLVisionBlock(nn.Module): + def __init__(self, config, attn_implementation: str = "sdpa") -> None: super().__init__() self.norm1 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) self.norm2 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) - self.attn = QWEN2_5_VL_VISION_ATTENTION_CLASSES[attn_implementation]( - config.hidden_size, num_heads=config.num_heads - ) + self.attn = QWEN2_5_VL_VISION_ATTENTION_CLASSES[attn_implementation](config.hidden_size, + num_heads=config.num_heads) self.mlp = Qwen2_5_VLMLP(config, bias=True) def forward(self, hidden_states, cu_seqlens, rotary_pos_emb) -> torch.Tensor: @@ -327,9 +331,7 @@ class Qwen2_5_VisionTransformerPretrainedModel(nn.Module): super().__init__() config_torch_dtype = getattr(config, "torch_dtype", None) - self.dtype = _resolve_torch_dtype( - config_torch_dtype if config_torch_dtype is not None else parent_torch_dtype - ) + self.dtype = _resolve_torch_dtype(config_torch_dtype if config_torch_dtype is not None else parent_torch_dtype) self.spatial_merge_size = config.spatial_merge_size self.patch_size = config.patch_size @@ -348,8 +350,7 @@ class Qwen2_5_VisionTransformerPretrainedModel(nn.Module): self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2) self.blocks = nn.ModuleList( - [Qwen2_5_VLVisionBlock(config, config._attn_implementation) for _ in range(config.depth)] - ) + [Qwen2_5_VLVisionBlock(config, config._attn_implementation) for _ in range(config.depth)]) self.merger = Qwen2_5_VLPatchMerger( dim=config.out_hidden_size, context_dim=config.hidden_size, @@ -472,9 +473,8 @@ class Qwen2_5_VisionTransformerPretrainedModel(nn.Module): for layer_num, blk in enumerate(self.blocks): cu_seqlens_now = cu_seqlens if layer_num in self.fullatt_block_indexes else cu_window_seqlens if self.gradient_checkpointing and self.training: - hidden_states = self._gradient_checkpointing_func( - blk.__call__, hidden_states, cu_seqlens_now, rotary_pos_emb - ) + hidden_states = self._gradient_checkpointing_func(blk.__call__, hidden_states, cu_seqlens_now, + rotary_pos_emb) else: hidden_states = blk( hidden_states, @@ -502,6 +502,7 @@ def _compute_default_rope_parameters(config, device=None, seq_len=None, **kwargs class Qwen2_5_VLRotaryEmbedding(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, device=None): super().__init__() # BC: "rope_type" was originally "type" @@ -576,6 +577,7 @@ class Qwen2_5_VLRotaryEmbedding(nn.Module): class Qwen2MLP(nn.Module): + def __init__(self, config): super().__init__() self.config = config @@ -662,8 +664,7 @@ class Qwen2_5_VLAttention(nn.Module): logger.warning_once( f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " - "when creating this class." - ) + "when creating this class.") self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads @@ -675,10 +676,8 @@ class Qwen2_5_VLAttention(nn.Module): self.rope_scaling = config.rope_scaling if (self.head_dim * self.num_heads) != self.hidden_size: - raise ValueError( - f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" - f" and `num_heads`: {self.num_heads})." - ) + raise ValueError(f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" + f" and `num_heads`: {self.num_heads}).") self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True) self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) @@ -711,9 +710,8 @@ class Qwen2_5_VLAttention(nn.Module): value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) cos, sin = position_embeddings - query_states, key_states = apply_multimodal_rotary_pos_emb( - query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] - ) + query_states, key_states = apply_multimodal_rotary_pos_emb(query_states, key_states, cos, sin, + self.rope_scaling["mrope_section"]) if past_key_value is not None: cache_kwargs = { @@ -730,7 +728,7 @@ class Qwen2_5_VLAttention(nn.Module): attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) if attention_mask is not None: # no matter the length, we just slice it - causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + causal_mask = attention_mask[:, :, :, :key_states.shape[-2]] attn_weights = attn_weights + causal_mask # Fix precision issues in Qwen2-VL float16 inference @@ -744,10 +742,8 @@ class Qwen2_5_VLAttention(nn.Module): attn_output = torch.matmul(attn_weights, value_states) if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): - raise ValueError( - f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" - f" {attn_output.size()}" - ) + raise ValueError(f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}") attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, -1) @@ -803,9 +799,8 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): # Because the input can be padded, the absolute sequence length depends on the max position id. cos, sin = position_embeddings - query_states, key_states = apply_multimodal_rotary_pos_emb( - query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] - ) + query_states, key_states = apply_multimodal_rotary_pos_emb(query_states, key_states, cos, sin, + self.rope_scaling["mrope_section"]) if past_key_value is not None: cache_kwargs = { @@ -836,8 +831,7 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): logger.warning_once( f"The input hidden states seems to be silently casted in float32, this might be related to" f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" - f" {target_dtype}." - ) + f" {target_dtype}.") query_states = query_states.to(target_dtype) key_states = key_states.to(target_dtype) @@ -848,11 +842,8 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) - if ( - self.config.use_sliding_window - and getattr(self.config, "sliding_window", None) is not None - and self.layer_idx >= self.config.max_window_layers - ): + if (self.config.use_sliding_window and getattr(self.config, "sliding_window", None) is not None + and self.layer_idx >= self.config.max_window_layers): sliding_window = self.config.sliding_window else: sliding_window = None @@ -929,9 +920,8 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) cos, sin = position_embeddings - query_states, key_states = apply_multimodal_rotary_pos_emb( - query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] - ) + query_states, key_states = apply_multimodal_rotary_pos_emb(query_states, key_states, cos, sin, + self.rope_scaling["mrope_section"]) if past_key_value is not None: cache_kwargs = { @@ -946,7 +936,7 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): causal_mask = attention_mask if attention_mask is not None and attention_mask.ndim == 4: # no matter the length, we just slice it - causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + causal_mask = attention_mask[:, :, :, :key_states.shape[-2]] # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, # Reference: https://github.com/pytorch/pytorch/issues/112577. @@ -985,7 +975,6 @@ QWEN2_5_VL_ATTENTION_CLASSES = { "sdpa": Qwen2_5_VLSdpaAttention, } - # If FlashAttention2 is not available, transparently fall back to SDPA. if not is_flash_attn_2_available(): QWEN2_5_VL_ATTENTION_CLASSES["flash_attention_2"] = Qwen2_5_VLSdpaAttention @@ -999,8 +988,10 @@ else: lazy_import_flash_attention("flash_attention_2") except (ImportError, ValueError): QWEN2_5_VL_ATTENTION_CLASSES["flash_attention_2"] = Qwen2_5_VLSdpaAttention - + + class Qwen2_5_VLDecoderLayer(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size @@ -1008,8 +999,7 @@ class Qwen2_5_VLDecoderLayer(nn.Module): if config.use_sliding_window and config._attn_implementation != "flash_attention_2": logger.warning_once( f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " - "unexpected results may be encountered." - ) + "unexpected results may be encountered.") self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx) self.mlp = Qwen2MLP(config) self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -1072,17 +1062,18 @@ class Qwen2_5_VLDecoderLayer(nn.Module): hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states - outputs = (hidden_states,) + outputs = (hidden_states, ) if output_attentions: - outputs += (self_attn_weights,) + outputs += (self_attn_weights, ) if use_cache: - outputs += (present_key_value,) + outputs += (present_key_value, ) return outputs class Qwen2_5_VLModel(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig): super().__init__() self.config = config @@ -1091,8 +1082,7 @@ class Qwen2_5_VLModel(nn.Module): self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( - [Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] - ) + [Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) self._attn_implementation = config._attn_implementation self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) @@ -1128,9 +1118,8 @@ class Qwen2_5_VLModel(nn.Module): ) -> Union[Tuple, BaseModelOutputWithPast]: assert not self.gradient_checkpointing output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else self.config.output_hidden_states) use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict @@ -1141,8 +1130,7 @@ class Qwen2_5_VLModel(nn.Module): if self.gradient_checkpointing and self.training: if use_cache: logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...") use_cache = False # torch.jit.trace() doesn't support cache objects in the output @@ -1191,18 +1179,18 @@ class Qwen2_5_VLModel(nn.Module): if self.cp_mesh is not None: seq_range = self._seq_range(seqlen) position_embeddings = ( - position_embeddings[0][:, :, seq_range[0] : seq_range[1], :], - position_embeddings[1][:, :, seq_range[0] : seq_range[1], :], + position_embeddings[0][:, :, seq_range[0]:seq_range[1], :], + position_embeddings[1][:, :, seq_range[0]:seq_range[1], :], ) - hidden_states = hidden_states[:, seq_range[0] : seq_range[1], :] - position_ids = position_ids[:, :, seq_range[0] : seq_range[1]] - cache_position = cache_position[seq_range[0] : seq_range[1]] - causal_mask = causal_mask[seq_range[0] : seq_range[1]] + hidden_states = hidden_states[:, seq_range[0]:seq_range[1], :] + position_ids = position_ids[:, :, seq_range[0]:seq_range[1]] + cache_position = cache_position[seq_range[0]:seq_range[1]] + causal_mask = causal_mask[seq_range[0]:seq_range[1]] assert past_key_values is None, "not support cp with past_key_values" for decoder_layer in self.layers: if output_hidden_states: - all_hidden_states += (hidden_states,) + all_hidden_states += (hidden_states, ) layer_id += 1 if self.gradient_checkpointing and self.training: raise NotImplementedError("Gradient checkpointing is not implemented here") @@ -1236,13 +1224,13 @@ class Qwen2_5_VLModel(nn.Module): next_decoder_cache = layer_outputs[2 if output_attentions else 1] if output_attentions: - all_self_attns += (layer_outputs[1],) + all_self_attns += (layer_outputs[1], ) hidden_states = self.norm(hidden_states) # add hidden states from the last decoder layer if output_hidden_states: - all_hidden_states += (hidden_states,) + all_hidden_states += (hidden_states, ) next_cache = next_decoder_cache if use_cache else None @@ -1279,8 +1267,7 @@ class Qwen2_5_VLModel(nn.Module): raise ValueError( "You are attempting to perform batched generation with padding_side='right'" " this may lead to unexpected behaviour for Flash Attention version of Qwen2_5_VL. Make sure to " - " call `tokenizer.padding_side = 'left'` before tokenizing the input. " - ) + " call `tokenizer.padding_side = 'left'` before tokenizing the input. ") if attention_mask is not None and 0.0 in attention_mask: return attention_mask return None @@ -1293,17 +1280,14 @@ class Qwen2_5_VLModel(nn.Module): using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache) # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if ( - self.config._attn_implementation == "sdpa" - and not (using_static_cache or using_sliding_window_cache) - and not output_attentions - ): + if (self.config._attn_implementation == "sdpa" and not (using_static_cache or using_sliding_window_cache) + and not output_attentions): if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - sliding_window=self.config.sliding_window, - is_training=self.training, + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + sliding_window=self.config.sliding_window, + is_training=self.training, ): return None @@ -1315,11 +1299,8 @@ class Qwen2_5_VLModel(nn.Module): target_length = past_key_values.get_max_cache_shape() # DynamicCache or no cache else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) + target_length = (attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) else past_seen_tokens + sequence_length + 1) # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( @@ -1334,12 +1315,8 @@ class Qwen2_5_VLModel(nn.Module): past_key_values=past_key_values, ) - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type in ["cuda", "xpu"] - and not output_attentions - ): + if (self.config._attn_implementation == "sdpa" and attention_mask is not None + and attention_mask.device.type in ["cuda", "xpu"] and not output_attentions): # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. # Details: https://github.com/pytorch/pytorch/issues/110213 @@ -1399,9 +1376,8 @@ class Qwen2_5_VLModel(nn.Module): # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also # the check is needed to verify is current checkpoint was trained with sliding window or not if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length: - sliding_attend_mask = torch.arange(target_length, device=device) <= ( - cache_position.reshape(-1, 1) - config.sliding_window - ) + sliding_attend_mask = torch.arange(target_length, device=device) <= (cache_position.reshape(-1, 1) - + config.sliding_window) diagonal_attend_mask.bitwise_or_(sliding_attend_mask) causal_mask *= diagonal_attend_mask causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) @@ -1411,12 +1387,10 @@ class Qwen2_5_VLModel(nn.Module): attention_mask = attention_mask[:, :target_length] mask_length = attention_mask.shape[-1] padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( - causal_mask.device - ) + causal_mask.device) padding_mask = padding_mask == 0 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) + padding_mask, min_dtype) return causal_mask @@ -1644,8 +1618,7 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module): st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 llm_pos_ids_list.append( - torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + st_idx - ) + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + st_idx) range_tensor = torch.arange(llm_grid_t).view(-1, 1) expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) @@ -1678,11 +1651,8 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module): max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0] mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] else: - position_ids = ( - torch.arange(input_ids.shape[1], device=input_ids.device) - .view(1, 1, -1) - .expand(3, input_ids.shape[0], -1) - ) + position_ids = (torch.arange(input_ids.shape[1], + device=input_ids.device).view(1, 1, -1).expand(3, input_ids.shape[0], -1)) mrope_position_deltas = torch.zeros( [input_ids.shape[0], 1], device=input_ids.device, @@ -1752,9 +1722,8 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module): ```""" output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else self.config.output_hidden_states) return_dict = return_dict if return_dict is not None else self.config.use_return_dict if inputs_embeds is None: @@ -1802,22 +1771,17 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module): video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) if is_inputs_embeds_dtensor: - inputs_embeds = ( - DTensor.from_local(inputs_embeds, device_mesh=target_device_mesh) - .redistribute(placements=target_placements) - .to_local() - ) + inputs_embeds = (DTensor.from_local( + inputs_embeds, + device_mesh=target_device_mesh).redistribute(placements=target_placements).to_local()) if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) # if we get 4D attention mask we cannot calculate rope deltas anymore. if position_ids is None and (attention_mask is None or attention_mask.ndim == 2): # calculate RoPE index once per generation in the pre-fill stage only - if ( - (cache_position is not None and cache_position[0] == 0) - or self.rope_deltas is None - or (past_key_values is None or past_key_values.get_seq_length() == 0) - ): + if ((cache_position is not None and cache_position[0] == 0) or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0)): position_ids, rope_deltas = self.get_rope_index( input_ids, image_grid_thw, @@ -1829,9 +1793,8 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module): # then use the prev pre-calculated rope-deltas to get the correct position ids else: batch_size, seq_length, _ = inputs_embeds.shape - delta = ( - (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) if cache_position is not None else 0 - ) + delta = ((cache_position[0] + + self.rope_deltas).to(inputs_embeds.device) if cache_position is not None else 0) position_ids = torch.arange(seq_length, device=inputs_embeds.device) position_ids = position_ids.view(1, -1).expand(batch_size, -1) if cache_position is not None: # otherwise `deltas` is an int `0` @@ -1871,8 +1834,8 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module): loss = loss_fct(shift_logits, shift_labels) if not return_dict: - output = (logits,) # + outputs[1:] - return (loss,) + output if loss is not None else output + output = (logits, ) # + outputs[1:] + return (loss, ) + output if loss is not None else output return Qwen2_5_VLCausalLMOutputWithPast( loss=loss, @@ -1910,11 +1873,10 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module): # generate the first token for each sequence. Later use the generated Input ids for continuation. if past_key_values is not None: if inputs_embeds is not None and input_ids.shape[1] == 0: # Exception 4 - inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :] + inputs_embeds = inputs_embeds[:, -cache_position.shape[0]:] elif inputs_embeds is not None or ( # Exception 1 - is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1] - ): # Exception 3 - input_ids = input_ids[:, -cache_position.shape[0] :] + is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1]): # Exception 3 + input_ids = input_ids[:, -cache_position.shape[0]:] elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2) input_ids = input_ids[:, cache_position] @@ -1948,20 +1910,18 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module): past_key_values=past_key_values, ) - model_inputs.update( - { - "position_ids": position_ids, - "past_key_values": past_key_values, - "use_cache": use_cache, - "attention_mask": attention_mask, - "pixel_values": pixel_values, - "pixel_values_videos": pixel_values_videos, - "image_grid_thw": image_grid_thw, - "video_grid_thw": video_grid_thw, - "cache_position": cache_position, - "second_per_grid_ts": second_per_grid_ts, - } - ) + model_inputs.update({ + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": use_cache, + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "pixel_values_videos": pixel_values_videos, + "image_grid_thw": image_grid_thw, + "video_grid_thw": video_grid_thw, + "cache_position": cache_position, + "second_per_grid_ts": second_per_grid_ts, + }) return model_inputs @@ -2097,8 +2057,7 @@ def get_rope_index( st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 llm_pos_ids_list.append( - torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + st_idx - ) + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + st_idx) range_tensor = torch.arange(llm_grid_t, device=input_ids.device).view(-1, 1) expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) @@ -2108,18 +2067,12 @@ def get_rope_index( time_tensor_long = time_tensor.long() t_index = time_tensor_long.flatten() - h_index = ( - torch.arange(llm_grid_h, device=input_ids.device) - .view(1, -1, 1) - .expand(llm_grid_t, -1, llm_grid_w) - .flatten() - ) - w_index = ( - torch.arange(llm_grid_w, device=input_ids.device) - .view(1, 1, -1) - .expand(llm_grid_t, llm_grid_h, -1) - .flatten() - ) + h_index = (torch.arange(llm_grid_h, device=input_ids.device).view(1, -1, + 1).expand(llm_grid_t, -1, + llm_grid_w).flatten()) + w_index = (torch.arange(llm_grid_w, + device=input_ids.device).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, + -1).flatten()) llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) st = ed + llm_grid_t * llm_grid_h * llm_grid_w @@ -2127,8 +2080,7 @@ def get_rope_index( st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 text_len = len(input_tokens) - st llm_pos_ids_list.append( - torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + st_idx - ) + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + st_idx) llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) @@ -2143,11 +2095,8 @@ def get_rope_index( max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0] mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] else: - position_ids = ( - torch.arange(input_ids.shape[1], device=input_ids.device) - .view(1, 1, -1) - .expand(3, input_ids.shape[0], -1) - ) + position_ids = (torch.arange(input_ids.shape[1], + device=input_ids.device).view(1, 1, -1).expand(3, input_ids.shape[0], -1)) mrope_position_deltas = torch.zeros( [input_ids.shape[0], 1], device=input_ids.device, diff --git a/fastvideo/models/encoders/qwen3.py b/fastvideo/models/encoders/qwen3.py index e664463ad..d09700cfe 100644 --- a/fastvideo/models/encoders/qwen3.py +++ b/fastvideo/models/encoders/qwen3.py @@ -56,9 +56,7 @@ class Qwen3MLP(nn.Module): prefix=f"{prefix}.down_proj", ) if hidden_act != "silu": - raise ValueError( - f"Unsupported activation: {hidden_act}. Only silu is supported." - ) + raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported.") self.act_fn = SiluAndMul() def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -100,9 +98,7 @@ class Qwen3Attention(nn.Module): assert tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = getattr( - config, "head_dim", self.hidden_size // self.total_num_heads - ) + self.head_dim = getattr(config, "head_dim", self.hidden_size // self.total_num_heads) self.rotary_dim = self.head_dim self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim @@ -225,9 +221,7 @@ class Qwen3DecoderLayer(nn.Module): config=config, hidden_size=self.hidden_size, num_heads=config.num_attention_heads, - num_kv_heads=getattr( - config, "num_key_value_heads", config.num_attention_heads - ), + num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads), rope_theta=rope_theta, rope_scaling=rope_scaling, max_position_embeddings=max_position_embeddings, @@ -244,9 +238,7 @@ class Qwen3DecoderLayer(nn.Module): prefix=f"{prefix}.mlp", ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, @@ -306,16 +298,13 @@ class Qwen3ForCausalLM(TextEncoder): quant_config=self.quant_config, ) - self.layers = nn.ModuleList( - [ - Qwen3DecoderLayer( - config=config, - quant_config=self.quant_config, - prefix=f"{config.prefix}.layers.{i}", - ) - for i in range(config.num_hidden_layers) - ] - ) + self.layers = nn.ModuleList([ + Qwen3DecoderLayer( + config=config, + quant_config=self.quant_config, + prefix=f"{config.prefix}.layers.{i}", + ) for i in range(config.num_hidden_layers) + ]) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -351,11 +340,8 @@ class Qwen3ForCausalLM(TextEncoder): output_hidden_states: bool | None = None, **kwargs: Any, ) -> BaseEncoderOutput: - output_hidden_states = ( - output_hidden_states - if output_hidden_states is not None - else self.config.output_hidden_states - ) + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else self.config.output_hidden_states) if inputs_embeds is not None: hidden_states = inputs_embeds @@ -379,17 +365,11 @@ class Qwen3ForCausalLM(TextEncoder): dtype=torch.long, ).unsqueeze(0).expand(hidden_states.shape[0], -1) - all_hidden_states: tuple[Any, ...] | None = ( - () if output_hidden_states else None - ) + all_hidden_states: tuple[Any, ...] | None = (() if output_hidden_states else None) for layer in self.layers: if all_hidden_states is not None: - all_hidden_states += ( - (hidden_states,) - if residual is None - else (hidden_states + residual,) - ) + all_hidden_states += ((hidden_states, ) if residual is None else (hidden_states + residual, )) hidden_states, residual = layer( position_ids, hidden_states, @@ -400,16 +380,14 @@ class Qwen3ForCausalLM(TextEncoder): hidden_states, _ = self.norm(hidden_states, residual) if all_hidden_states is not None: - all_hidden_states += (hidden_states,) + all_hidden_states += (hidden_states, ) return BaseEncoderOutput( last_hidden_state=hidden_states, hidden_states=all_hidden_states, ) - def load_weights( - self, weights: Iterable[tuple[str, torch.Tensor]] - ) -> set[str]: + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() stacked_params_mapping = self.config.arch_config.stacked_params_mapping @@ -417,12 +395,9 @@ class Qwen3ForCausalLM(TextEncoder): # or after every split source projection has been loaded. Include # auxiliary quantization parameters (for example scale_weight) rather # than limiting completeness checks to weight/bias tensors. - expected_stacked_shards = { - (name, shard_id) - for name in params_dict - for param_name, _, shard_id in stacked_params_mapping - if param_name in name - } + expected_stacked_shards = {(name, shard_id) + for name in params_dict + for param_name, _, shard_id in stacked_params_mapping if param_name in name} fused_param_names = {name for name, _ in expected_stacked_shards} loaded_stacked_shards: set[tuple[str, str | int]] = set() loaded_fused_params: set[str] = set() @@ -437,9 +412,7 @@ class Qwen3ForCausalLM(TextEncoder): continue if "scale" in name: - kv_scale_name: str | None = maybe_remap_kv_scale_name( - name, params_dict - ) + kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict) if kv_scale_name is None: continue name = kv_scale_name @@ -486,24 +459,16 @@ class Qwen3ForCausalLM(TextEncoder): if name in fused_param_names: loaded_fused_params.add(name) - required_split_shards = { - (name, shard_id) - for name, shard_id in expected_stacked_shards - if name not in loaded_fused_params - } + required_split_shards = {(name, shard_id) + for name, shard_id in expected_stacked_shards if name not in loaded_fused_params} missing_stacked_shards = required_split_shards - loaded_stacked_shards if missing_stacked_shards: - formatted_missing = ", ".join( - f"{name}[{shard_id}]" - for name, shard_id in sorted( - missing_stacked_shards, - key=lambda item: (item[0], str(item[1])), - ) - ) - raise ValueError( - "Missing required stacked checkpoint shards: " - f"{formatted_missing}" - ) + formatted_missing = ", ".join(f"{name}[{shard_id}]" for name, shard_id in sorted( + missing_stacked_shards, + key=lambda item: (item[0], str(item[1])), + )) + raise ValueError("Missing required stacked checkpoint shards: " + f"{formatted_missing}") return loaded_params diff --git a/fastvideo/models/encoders/reason1.py b/fastvideo/models/encoders/reason1.py index fc3f6516c..51f9f711c 100644 --- a/fastvideo/models/encoders/reason1.py +++ b/fastvideo/models/encoders/reason1.py @@ -33,8 +33,6 @@ class _WeightsSource: allow_patterns_overrides: list[str] | None = None - - class Reason1TextEncoder(TextEncoder): """Reason1 (Qwen2.5-VL) text encoder.""" @@ -42,19 +40,19 @@ class Reason1TextEncoder(TextEncoder): AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, ) - + def __init__(self, config: Reason1Config, prefix: str = "", checkpoint_path: str | None = None): super().__init__(config) - + self.prefix = prefix self.quant_config = None # For future quantization support - + self.embedding_concat_strategy = config.arch_config.embedding_concat_strategy self.n_layers_per_group = config.arch_config.n_layers_per_group self.num_embedding_padding_tokens = config.arch_config.num_embedding_padding_tokens - + config_path = checkpoint_path if checkpoint_path else config.tokenizer_type - + logger.info("Initializing Reason1TextEncoder (Qwen2.5-VL) from %s", config_path) try: from transformers import AutoConfig as HFAutoConfig @@ -63,30 +61,29 @@ class Reason1TextEncoder(TextEncoder): trust_remote_code=True, ) except Exception as e: - logger.warning("Failed to load HF config from %s (%s). Using default Qwen2.5-VL-7B config.", - config_path, e) + logger.warning("Failed to load HF config from %s (%s). Using default Qwen2.5-VL-7B config.", config_path, e) hf_config = Qwen2_5_VLConfig( - hidden_size=3584, - intermediate_size=18944, - max_window_layers=28, - num_attention_heads=28, - num_hidden_layers=28, - num_key_value_heads=4, - tie_word_embeddings=False, - vocab_size=152064, + hidden_size=3584, + intermediate_size=18944, + max_window_layers=28, + num_attention_heads=28, + num_hidden_layers=28, + num_key_value_heads=4, + tie_word_embeddings=False, + vocab_size=152064, ) - + hf_config.output_hidden_states = True - + if hasattr(config.arch_config, '_attn_implementation') and config.arch_config._attn_implementation: hf_config._attn_implementation = config.arch_config._attn_implementation else: hf_config._attn_implementation = "flash_attention_2" logger.info("Reason1 attention implementation: %s", getattr(hf_config, "_attn_implementation", None)) - + with torch.device("meta"): self.model = Qwen2_5_VLForConditionalGenerationSimple(hf_config) - + self.processor = AutoProcessor.from_pretrained( config_path, trust_remote_code=True, @@ -94,18 +91,16 @@ class Reason1TextEncoder(TextEncoder): weights_override = os.getenv("FASTVIDEO_REASON1_WEIGHTS_PATH") if weights_override: - self.secondary_weights = ( - _WeightsSource( - model_or_path=weights_override, - prefix="", - fall_back_to_pt=True, - allow_patterns_overrides=None, - ), - ) + self.secondary_weights = (_WeightsSource( + model_or_path=weights_override, + prefix="", + fall_back_to_pt=True, + allow_patterns_overrides=None, + ), ) logger.info("Reason1TextEncoder: overlaying weights from %s", weights_override) - + self._weights_loaded = False - + def forward( self, input_ids: torch.Tensor | None, @@ -128,16 +123,16 @@ class Reason1TextEncoder(TextEncoder): image_grid_thw=kwargs.get('image_grid_thw', None), video_grid_thw=kwargs.get('video_grid_thw', None), ) - + hidden_states = outputs.hidden_states last_hidden_state = hidden_states[-1] - + return BaseEncoderOutput( last_hidden_state=last_hidden_state, hidden_states=hidden_states if output_hidden_states else None, attention_mask=None, ) - + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: first_weight = None weights_list = [] @@ -147,11 +142,11 @@ class Reason1TextEncoder(TextEncoder): self.model = self.model.to_empty(device=weight.device) self.model.init_weights(buffer_device=weight.device) weights_list.append((name, weight)) - + params_dict = dict(self.model.named_parameters()) loaded_params: set[str] = set() skipped_weights = {"lm_head": 0, "visual": 0, "decoder": 0} - + for name, loaded_weight in weights_list: if "lm_head" in name: skipped_weights["lm_head"] += 1 @@ -162,7 +157,7 @@ class Reason1TextEncoder(TextEncoder): if "decoder" in name: skipped_weights["decoder"] += 1 continue - + # Handle stacked params mapping (for quantized models) for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping: if weight_name not in name: @@ -170,10 +165,10 @@ class Reason1TextEncoder(TextEncoder): name = name.replace(weight_name, param_name) if name.endswith(".bias") and name not in params_dict: continue - + if name not in params_dict: continue - + param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) @@ -183,10 +178,10 @@ class Reason1TextEncoder(TextEncoder): else: if name.endswith(".bias") and name not in params_dict: continue - + if name not in params_dict: continue - + param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) @@ -194,15 +189,15 @@ class Reason1TextEncoder(TextEncoder): loaded_params.add(param_name_with_prefix) if first_weight is not None: self.model = self.model.to(first_weight.device) - - all_params = set(f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}" - for name in params_dict.keys()) + + all_params = set(f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}" + for name in params_dict.keys()) loaded_params.update(all_params) - + # Mark weights as loaded self._weights_loaded = True return loaded_params - + def compute_text_embeddings_online( self, data_batch: dict[str, list[str]], @@ -210,7 +205,7 @@ class Reason1TextEncoder(TextEncoder): ) -> torch.Tensor: prompts = data_batch[input_caption_key] return self.compute_text_embeddings(prompts) - + def compute_text_embeddings( self, prompts: list[str], @@ -218,7 +213,7 @@ class Reason1TextEncoder(TextEncoder): ) -> torch.Tensor: """Compute embeddings for a list of prompts.""" input_ids_batch = [] - + tok = getattr(self.processor, "tokenizer", None) if tok is None: raise RuntimeError("Reason1TextEncoder requires processor.tokenizer") @@ -229,29 +224,26 @@ class Reason1TextEncoder(TextEncoder): pad_id = getattr(self.model.config, "pad_token_id", None) if pad_id is None: pad_id = 0 - + for prompt in prompts: conversations = [ { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a helpful assistant who will provide prompts to an image generator.", - } - ], + "role": + "system", + "content": [{ + "type": "text", + "text": "You are a helpful assistant who will provide prompts to an image generator.", + }], }, { "role": "user", - "content": [ - { - "type": "text", - "text": prompt, - } - ], + "content": [{ + "type": "text", + "text": prompt, + }], }, ] - + try: tokenizer_output = tok.apply_chat_template( conversations, @@ -261,10 +253,10 @@ class Reason1TextEncoder(TextEncoder): ) except TypeError: tokenizer_output = tok.apply_chat_template( - conversations, - tokenize=True, - add_generation_prompt=False, - ) + conversations, + tokenize=True, + add_generation_prompt=False, + ) if isinstance(tokenizer_output, dict) and "input_ids" in tokenizer_output: input_ids = tokenizer_output["input_ids"] @@ -274,23 +266,20 @@ class Reason1TextEncoder(TextEncoder): input_ids = tokenizer_output if hasattr(input_ids, "tolist"): input_ids = input_ids.tolist() - if isinstance(input_ids, list) and len(input_ids) == 1 and isinstance( - input_ids[0], list): + if isinstance(input_ids, list) and len(input_ids) == 1 and isinstance(input_ids[0], list): input_ids = input_ids[0] if not isinstance(input_ids, list): - raise RuntimeError( - f"Unexpected chat_template output type: {type(tokenizer_output)}" - ) - + raise RuntimeError(f"Unexpected chat_template output type: {type(tokenizer_output)}") + if self.num_embedding_padding_tokens > len(input_ids): pad_len = self.num_embedding_padding_tokens - len(input_ids) input_ids = input_ids + [pad_id] * pad_len else: input_ids = input_ids[:self.num_embedding_padding_tokens] - + input_ids = torch.LongTensor(input_ids).to(device=device) input_ids_batch.append(input_ids) - + input_ids_batch = torch.stack(input_ids_batch, dim=0) # Cosmos2.5 alignment: keep attention_mask=None. @@ -301,7 +290,7 @@ class Reason1TextEncoder(TextEncoder): embed_device = None if embed_device is not None and embed_device != target_device: self.model = self.model.to(target_device) - + with torch.no_grad(): position_ids, _ = get_rope_index( self.model.config, @@ -312,7 +301,7 @@ class Reason1TextEncoder(TextEncoder): attention_mask=None, ) position_ids = position_ids.to(target_device) - + outputs = self.model.model( input_ids=input_ids_batch, position_ids=position_ids, @@ -322,12 +311,12 @@ class Reason1TextEncoder(TextEncoder): use_cache=False, ) hidden_states = outputs.hidden_states - + normalized_hidden_states = [] for layer_idx in range(1, len(hidden_states)): normalized_state = self._mean_normalize(hidden_states[layer_idx]) normalized_hidden_states.append(normalized_state) - + if self.embedding_concat_strategy == "full_concat": text_embeddings = torch.cat(normalized_hidden_states, dim=-1) elif self.embedding_concat_strategy == "mean_pooling": @@ -335,21 +324,19 @@ class Reason1TextEncoder(TextEncoder): elif self.embedding_concat_strategy == "pool_every_n_layers_and_concat": pooled_embeddings = [] for i in range(0, len(normalized_hidden_states), self.n_layers_per_group): - group = normalized_hidden_states[i : i + self.n_layers_per_group] + group = normalized_hidden_states[i:i + self.n_layers_per_group] pooled = torch.stack(group).mean(dim=0) pooled_embeddings.append(pooled) text_embeddings = torch.cat(pooled_embeddings, dim=-1) else: - raise ValueError( - f"Unknown embedding_concat_strategy: {self.embedding_concat_strategy}" - ) - + raise ValueError(f"Unknown embedding_concat_strategy: {self.embedding_concat_strategy}") + return text_embeddings + @staticmethod def _mean_normalize(tensor: torch.Tensor) -> torch.Tensor: - return (tensor - tensor.mean(dim=-1, keepdim=True)) / ( - tensor.std(dim=-1, keepdim=True) + 1e-8 - ) + return (tensor - tensor.mean(dim=-1, keepdim=True)) / (tensor.std(dim=-1, keepdim=True) + 1e-8) + # Entry point for model registry EntryClass = Reason1TextEncoder diff --git a/fastvideo/models/encoders/siglip.py b/fastvideo/models/encoders/siglip.py index 2842102a4..0ac5cf506 100644 --- a/fastvideo/models/encoders/siglip.py +++ b/fastvideo/models/encoders/siglip.py @@ -16,8 +16,7 @@ from fastvideo.configs.models.encoders import BaseEncoderOutput from fastvideo.configs.models.encoders.siglip import SiglipVisionArchConfig, SiglipVisionConfig from fastvideo.distributed import divide, get_tp_world_size from fastvideo.layers.activation import get_act_fn -from fastvideo.layers.linear import (ColumnParallelLinear, QKVParallelLinear, - RowParallelLinear) +from fastvideo.layers.linear import (ColumnParallelLinear, QKVParallelLinear, RowParallelLinear) from fastvideo.layers.quantization import QuantizationConfig from fastvideo.logger import init_logger from fastvideo.models.encoders.base import ImageEncoder @@ -48,7 +47,7 @@ class SiglipVisionEmbeddings(nn.Module): ) # Integer division - with valid padding, edge pixels are ignored - self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_patches = (self.image_size // self.patch_size)**2 self.num_positions = self.num_patches self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) self.register_buffer( @@ -59,9 +58,7 @@ class SiglipVisionEmbeddings(nn.Module): def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: target_dtype = self.patch_embedding.weight.dtype - patch_embeds = self.patch_embedding( - pixel_values.to(dtype=target_dtype) - ) # shape = [*, embed_dim, grid, grid] + patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, embed_dim, grid, grid] embeddings = patch_embeds.flatten(2).transpose(1, 2) embeddings = embeddings + self.position_embedding(self.position_ids) return embeddings @@ -83,13 +80,11 @@ class SiglipAttention(nn.Module): self.embed_dim = arch_config.hidden_size self.num_heads = arch_config.num_attention_heads self.head_dim = self.embed_dim // self.num_heads - + if self.head_dim * self.num_heads != self.embed_dim: - raise ValueError( - f"embed_dim must be divisible by num_heads " - f"(got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})." - ) - + raise ValueError(f"embed_dim must be divisible by num_heads " + f"(got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads}).") + self.scale = self.head_dim**-0.5 if enable_scale else None self.dropout = arch_config.attention_dropout @@ -124,25 +119,17 @@ class SiglipAttention(nn.Module): """Input shape: Batch x Time x Channel""" qkv_states, _ = self.qkv_proj(hidden_states) query_states, key_states, value_states = qkv_states.chunk(3, dim=-1) - - query_states = query_states.reshape( - query_states.shape[0], query_states.shape[1], - self.num_heads_per_partition, self.head_dim - ) - key_states = key_states.reshape( - key_states.shape[0], key_states.shape[1], - self.num_heads_per_partition, self.head_dim - ) - value_states = value_states.reshape( - value_states.shape[0], value_states.shape[1], - self.num_heads_per_partition, self.head_dim - ) - + + query_states = query_states.reshape(query_states.shape[0], query_states.shape[1], self.num_heads_per_partition, + self.head_dim) + key_states = key_states.reshape(key_states.shape[0], key_states.shape[1], self.num_heads_per_partition, + self.head_dim) + value_states = value_states.reshape(value_states.shape[0], value_states.shape[1], self.num_heads_per_partition, + self.head_dim) + attn_output = self.attn(query_states, key_states, value_states) - attn_output = attn_output.reshape( - attn_output.shape[0], attn_output.shape[1], - self.num_heads_per_partition * self.head_dim - ) + attn_output = attn_output.reshape(attn_output.shape[0], attn_output.shape[1], + self.num_heads_per_partition * self.head_dim) attn_output, _ = self.out_proj(attn_output) return attn_output, None @@ -198,17 +185,13 @@ class SiglipEncoderLayer(nn.Module): quant_config=quant_config, prefix=f"{prefix}.self_attn", ) - self.layer_norm1 = nn.LayerNorm( - arch_config.hidden_size, eps=arch_config.layer_norm_eps - ) + self.layer_norm1 = nn.LayerNorm(arch_config.hidden_size, eps=arch_config.layer_norm_eps) self.mlp = SiglipMLP( arch_config, quant_config=quant_config, prefix=f"{prefix}.mlp", ) - self.layer_norm2 = nn.LayerNorm( - arch_config.hidden_size, eps=arch_config.layer_norm_eps - ) + self.layer_norm2 = nn.LayerNorm(arch_config.hidden_size, eps=arch_config.layer_norm_eps) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # SigLIP uses post-norm (like original ViT) @@ -244,7 +227,7 @@ class SiglipEncoder(nn.Module): num_hidden_layers = arch_config.num_hidden_layers else: num_hidden_layers = num_hidden_layers_override - + self.layers = nn.ModuleList([ SiglipEncoderLayer( arch_config=arch_config, @@ -252,8 +235,7 @@ class SiglipEncoder(nn.Module): is_causal=is_causal, quant_config=quant_config, prefix=f"{prefix}.layers.{layer_idx}", - ) - for layer_idx in range(num_hidden_layers) + ) for layer_idx in range(num_hidden_layers) ]) def forward( @@ -302,10 +284,8 @@ class SiglipVisionTransformer(nn.Module): num_hidden_layers = arch_config.num_hidden_layers if len(self.encoder.layers) > arch_config.num_hidden_layers: - raise ValueError( - f"The original encoder only has {num_hidden_layers} " - f"layers, but you requested {len(self.encoder.layers)} layers." - ) + raise ValueError(f"The original encoder only has {num_hidden_layers} " + f"layers, but you requested {len(self.encoder.layers)} layers.") # Post layer norm (applied to output) if require_post_norm is None: @@ -344,7 +324,7 @@ class SiglipVisionTransformer(nn.Module): class SiglipVisionModel(ImageEncoder): """SigLIP Vision Model for FastVideo.""" - + config_class = SiglipVisionConfig main_input_name = "pixel_values" packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} @@ -382,14 +362,13 @@ class SiglipVisionModel(ImageEncoder): # Skip projection layers if any if name.startswith("visual_projection"): continue - + # Skip head if any if "head" in name: continue # Post layernorm handling - if (name.startswith("vision_model.post_layernorm") - and self.vision_model.post_layernorm is None): + if (name.startswith("vision_model.post_layernorm") and self.vision_model.post_layernorm is None): continue # Omit layers when num_hidden_layers_override is set @@ -415,8 +394,9 @@ class SiglipVisionModel(ImageEncoder): weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) - + return loaded_params + # Entry point for model registry EntryClass = SiglipVisionModel diff --git a/fastvideo/models/encoders/stable_audio_conditioner.py b/fastvideo/models/encoders/stable_audio_conditioner.py index c4aba0105..c353ebce0 100644 --- a/fastvideo/models/encoders/stable_audio_conditioner.py +++ b/fastvideo/models/encoders/stable_audio_conditioner.py @@ -32,8 +32,7 @@ class _LearnedPositionalEmbedding(nn.Module): def _time_positional_embedding(dim: int, out_features: int) -> nn.Sequential: - return nn.Sequential(_LearnedPositionalEmbedding(dim), - nn.Linear(in_features=dim + 1, out_features=out_features)) + return nn.Sequential(_LearnedPositionalEmbedding(dim), nn.Linear(in_features=dim + 1, out_features=out_features)) class NumberEmbedder(nn.Module): @@ -59,8 +58,7 @@ class _Conditioner(nn.Module): super().__init__() self.dim = dim self.output_dim = output_dim - self.proj_out = (nn.Linear(dim, output_dim) if dim != output_dim or project_out - else nn.Identity()) + self.proj_out = (nn.Linear(dim, output_dim) if dim != output_dim or project_out else nn.Identity()) class T5Conditioner(_Conditioner): @@ -71,8 +69,11 @@ class T5Conditioner(_Conditioner): T5_MODEL_DIMS = {"t5-base": 768} - def __init__(self, output_dim: int, t5_model_name: str = "t5-base", - max_length: int = 128, dtype: str = "float16") -> None: + def __init__(self, + output_dim: int, + t5_model_name: str = "t5-base", + max_length: int = 128, + dtype: str = "float16") -> None: super().__init__(self.T5_MODEL_DIMS[t5_model_name], output_dim, project_out=False) from transformers import AutoTokenizer, T5EncoderModel self.max_length = max_length @@ -92,14 +93,16 @@ class T5Conditioner(_Conditioner): self.model = (T5EncoderModel.from_pretrained(t5_model_name).eval().requires_grad_(False).to(torch_dtype)) def forward(self, texts: list[str], device: torch.device | str) -> tuple[torch.Tensor, torch.Tensor]: - encoded = self.tokenizer(texts, truncation=True, max_length=self.max_length, - padding="max_length", return_tensors="pt") + encoded = self.tokenizer(texts, + truncation=True, + max_length=self.max_length, + padding="max_length", + return_tensors="pt") input_ids = encoded["input_ids"].to(device) attention_mask = encoded["attention_mask"].to(device).to(torch.bool) # Mirror official's `autocast(fp16)` wrap on T5 forward. with torch.no_grad(), torch.autocast(device_type="cuda", dtype=self._t5_dtype): - embeddings = self.model(input_ids=input_ids, - attention_mask=attention_mask)["last_hidden_state"] + embeddings = self.model(input_ids=input_ids, attention_mask=attention_mask)["last_hidden_state"] embeddings = self.proj_out(embeddings) * attention_mask.unsqueeze(-1).float() return embeddings, attention_mask @@ -149,8 +152,7 @@ class StableAudioMultiConditioner(nn.Module): max_length=scfg["max_length"], dtype=arch.t5_dtype) elif stype == "number": - sub[sid] = NumberConditioner(output_dim=arch.cond_dim, - min_val=scfg["min_val"], max_val=scfg["max_val"]) + sub[sid] = NumberConditioner(output_dim=arch.cond_dim, min_val=scfg["min_val"], max_val=scfg["max_val"]) else: raise ValueError(f"Unknown sub-conditioner type {stype!r} for id {sid!r}.") self.conditioners = nn.ModuleDict(sub) @@ -166,8 +168,8 @@ class StableAudioMultiConditioner(nn.Module): return out def get_conditioning_inputs( - self, cond: dict[str, tuple[torch.Tensor, torch.Tensor]] - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + self, cond: dict[str, tuple[torch.Tensor, + torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Pack conditioner outputs into the (cross_attn_cond, cross_attn_mask, global_embed) triple the DiT consumes. Order is driven by `cross_attention_cond_ids` / `global_cond_ids` @@ -183,7 +185,8 @@ class StableAudioMultiConditioner(nn.Module): return cross_attn_cond, cross_attn_mask, global_embed @classmethod - def from_official_state_dict(cls, state_dict: dict[str, torch.Tensor], + def from_official_state_dict(cls, + state_dict: dict[str, torch.Tensor], prefix: str = "conditioner.") -> "StableAudioMultiConditioner": """Load NumberConditioner weights from a raw `stable_audio_tools` monolithic state dict. Kept for tests / older checkpoints; @@ -200,13 +203,11 @@ class StableAudioMultiConditioner(nn.Module): if stripped in own_state: loaded[stripped] = v # T5 keys are intentionally absent from the checkpoint. - missing = [k for k in own_state.keys() if k not in loaded - and not k.startswith("conditioners.prompt.")] + missing = [k for k in own_state.keys() if k not in loaded and not k.startswith("conditioners.prompt.")] unexpected = [k for k in loaded.keys() if k not in own_state] if missing or unexpected: raise RuntimeError( - f"StableAudioMultiConditioner load mismatch — missing={missing[:5]} unexpected={unexpected[:5]}" - ) + f"StableAudioMultiConditioner load mismatch — missing={missing[:5]} unexpected={unexpected[:5]}") mc.load_state_dict(loaded, strict=False) return mc diff --git a/fastvideo/models/encoders/t5.py b/fastvideo/models/encoders/t5.py index e2f9b12da..e85e130dc 100644 --- a/fastvideo/models/encoders/t5.py +++ b/fastvideo/models/encoders/t5.py @@ -64,16 +64,11 @@ class AttentionMetadata: class T5DenseActDense(nn.Module): - def __init__( - self, config: T5Config, quant_config: QuantizationConfig | None = None - ): + + def __init__(self, config: T5Config, quant_config: QuantizationConfig | None = None): super().__init__() - self.wi = MergedColumnParallelLinear( - config.d_model, [config.d_ff], bias=False - ) - self.wo = RowParallelLinear( - config.d_ff, config.d_model, bias=False, quant_config=quant_config - ) + self.wi = MergedColumnParallelLinear(config.d_model, [config.d_ff], bias=False) + self.wo = RowParallelLinear(config.d_ff, config.d_model, bias=False, quant_config=quant_config) self.act = get_act_fn(config.dense_act_fn) def forward(self, hidden_states) -> torch.Tensor: @@ -84,21 +79,14 @@ class T5DenseActDense(nn.Module): class T5DenseGatedActDense(nn.Module): - def __init__( - self, config: T5Config, quant_config: QuantizationConfig | None = None - ): + + def __init__(self, config: T5Config, quant_config: QuantizationConfig | None = None): super().__init__() - self.wi_0 = MergedColumnParallelLinear( - config.d_model, [config.d_ff], bias=False, quant_config=quant_config - ) - self.wi_1 = MergedColumnParallelLinear( - config.d_model, [config.d_ff], bias=False, quant_config=quant_config - ) + self.wi_0 = MergedColumnParallelLinear(config.d_model, [config.d_ff], bias=False, quant_config=quant_config) + self.wi_1 = MergedColumnParallelLinear(config.d_model, [config.d_ff], bias=False, quant_config=quant_config) # Should not run in fp16 unless mixed-precision is used, # see https://github.com/huggingface/transformers/issues/20287. - self.wo = RowParallelLinear( - config.d_ff, config.d_model, bias=False, quant_config=quant_config - ) + self.wo = RowParallelLinear(config.d_ff, config.d_model, bias=False, quant_config=quant_config) self.act = get_act_fn(config.dense_act_fn) def forward(self, hidden_states) -> torch.Tensor: @@ -110,18 +98,13 @@ class T5DenseGatedActDense(nn.Module): class T5LayerFF(nn.Module): - def __init__( - self, config: T5Config, quant_config: QuantizationConfig | None = None - ): + + def __init__(self, config: T5Config, quant_config: QuantizationConfig | None = None): super().__init__() if config.is_gated_act: - self.DenseReluDense = T5DenseGatedActDense( - config, quant_config=quant_config - ) + self.DenseReluDense = T5DenseGatedActDense(config, quant_config=quant_config) else: - self.DenseReluDense = T5DenseActDense( - config, quant_config=quant_config - ) + self.DenseReluDense = T5DenseActDense(config, quant_config=quant_config) self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon) @@ -134,6 +117,7 @@ class T5LayerFF(nn.Module): # T5 has attn_bias and does not use softmax scaling class T5MultiHeadAttention(nn.Module): + def __init__(self) -> None: super().__init__() @@ -150,6 +134,7 @@ class T5MultiHeadAttention(nn.Module): class T5Attention(nn.Module): + def __init__( self, config: T5Config, @@ -163,12 +148,8 @@ class T5Attention(nn.Module): # Cross-attention has no relative pos encoding anyway self.is_decoder = attn_type == AttentionType.DECODER self.has_relative_attention_bias = has_relative_attention_bias - self.relative_attention_num_buckets = ( - config.relative_attention_num_buckets - ) - self.relative_attention_max_distance = ( - config.relative_attention_max_distance - ) + self.relative_attention_num_buckets = (config.relative_attention_num_buckets) + self.relative_attention_max_distance = (config.relative_attention_max_distance) self.d_model = config.d_model self.key_value_proj_dim = config.d_kv self.total_num_heads = self.total_num_kv_heads = config.num_heads @@ -211,9 +192,10 @@ class T5Attention(nn.Module): ) @staticmethod - def _relative_position_bucket( - relative_position, bidirectional=True, num_buckets=32, max_distance=128 - ) -> torch.Tensor: + def _relative_position_bucket(relative_position, + bidirectional=True, + num_buckets=32, + max_distance=128) -> torch.Tensor: """ Adapted from Mesh Tensorflow: https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 @@ -239,14 +221,10 @@ class T5Attention(nn.Module): relative_buckets = 0 if bidirectional: num_buckets //= 2 - relative_buckets += (relative_position > 0).to( - torch.long - ) * num_buckets + relative_buckets += (relative_position > 0).to(torch.long) * num_buckets relative_position = torch.abs(relative_position) else: - relative_position = -torch.min( - relative_position, torch.zeros_like(relative_position) - ) + relative_position = -torch.min(relative_position, torch.zeros_like(relative_position)) # now relative_position is in the range [0, inf) # half of the buckets are for exact increments in positions @@ -255,33 +233,23 @@ class T5Attention(nn.Module): # The other half of the buckets are for logarithmically bigger bins # in positions up to max_distance - relative_position_if_large = max_exact + ( - torch.log(relative_position.float() / max_exact) - / math.log(max_distance / max_exact) - * (num_buckets - max_exact) - ).to(torch.long) + relative_position_if_large = max_exact + (torch.log(relative_position.float() / max_exact) / + math.log(max_distance / max_exact) * + (num_buckets - max_exact)).to(torch.long) relative_position_if_large = torch.min( relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1), ) - relative_buckets += torch.where( - is_small, relative_position, relative_position_if_large - ) + relative_buckets += torch.where(is_small, relative_position, relative_position_if_large) return relative_buckets - def compute_bias( - self, query_length, key_length, device=None - ) -> torch.Tensor: + def compute_bias(self, query_length, key_length, device=None) -> torch.Tensor: """Compute binned relative position bias""" if device is None: device = self.relative_attention_bias.weight.device - context_position = torch.arange( - query_length, dtype=torch.long, device=device - )[:, None] - memory_position = torch.arange( - key_length, dtype=torch.long, device=device - )[None, :] + context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] + memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] # max_seq_len, nh relative_position = memory_position - context_position relative_position_bucket = self._relative_position_bucket( @@ -290,12 +258,8 @@ class T5Attention(nn.Module): num_buckets=self.relative_attention_num_buckets, max_distance=self.relative_attention_max_distance, ) - values = self.relative_attention_bias( - relative_position_bucket - ) # shape (query_length, key_length, num_heads) - x = values.permute([2, 0, 1]).unsqueeze( - 0 - ) # shape (1, num_heads, query_length, key_length) + values = self.relative_attention_bias(relative_position_bucket) # shape (query_length, key_length, num_heads) + x = values.permute([2, 0, 1]).unsqueeze(0) # shape (1, num_heads, query_length, key_length) return x def forward( @@ -323,9 +287,7 @@ class T5Attention(nn.Module): # The bias term is computed on longest sequence in batch. Biases # for shorter sequences are slices of the longest. assert self.attn_type == AttentionType.ENCODER - attn_bias = self.compute_bias(seq_len, seq_len).repeat( - num_seqs, 1, 1, 1 - ) + attn_bias = self.compute_bias(seq_len, seq_len).repeat(num_seqs, 1, 1, 1) attn_metadata.attn_bias = attn_bias else: # Encoder/Decoder Self-Attention Layer, attn bias already cached. @@ -334,27 +296,21 @@ class T5Attention(nn.Module): from fastvideo.platforms import current_platform if attention_mask is not None: - attention_mask = ( - attention_mask.view(bs, 1, 1, -1) - if attention_mask.ndim == 2 - else attention_mask.unsqueeze(1) - ) - mask_val = ( - -1e4 if current_platform.is_mps() else torch.finfo(q.dtype).min - ) + attention_mask = (attention_mask.view(bs, 1, 1, -1) + if attention_mask.ndim == 2 else attention_mask.unsqueeze(1)) + mask_val = (-1e4 if current_platform.is_mps() else torch.finfo(q.dtype).min) attn_bias.masked_fill_(attention_mask == 0, mask_val) if get_tp_world_size() > 1: rank = get_tp_rank() - attn_bias = attn_bias[ - :, rank * self.n_heads : (rank + 1) * self.n_heads, :, : - ] + attn_bias = attn_bias[:, rank * self.n_heads:(rank + 1) * self.n_heads, :, :] attn_output = self.attn(q, k, v, attn_bias) output, _ = self.o(attn_output) return output class T5LayerSelfAttention(nn.Module): + def __init__( self, config, @@ -365,9 +321,7 @@ class T5LayerSelfAttention(nn.Module): super().__init__() self.SelfAttention = T5Attention( config, - AttentionType.DECODER - if "decoder" in prefix - else AttentionType.ENCODER, + AttentionType.DECODER if "decoder" in prefix else AttentionType.ENCODER, has_relative_attention_bias=has_relative_attention_bias, quant_config=quant_config, prefix=f"{prefix}.SelfAttention", @@ -391,6 +345,7 @@ class T5LayerSelfAttention(nn.Module): class T5LayerCrossAttention(nn.Module): + def __init__( self, config, @@ -422,6 +377,7 @@ class T5LayerCrossAttention(nn.Module): class T5Block(nn.Module): + def __init__( self, config: T5Config, @@ -439,17 +395,14 @@ class T5Block(nn.Module): has_relative_attention_bias=has_relative_attention_bias, quant_config=quant_config, prefix=f"{prefix}.self_attn", - ) - ) + )) if self.is_decoder: - self.layer.append( - T5LayerCrossAttention( - config, - quant_config=quant_config, - prefix=f"{prefix}.cross_attn", - ) - ) + self.layer.append(T5LayerCrossAttention( + config, + quant_config=quant_config, + prefix=f"{prefix}.cross_attn", + )) self.layer.append(T5LayerFF(config, quant_config=quant_config)) @@ -465,9 +418,7 @@ class T5Block(nn.Module): attn_metadata=attn_metadata, ) if self.is_decoder: - hidden_states = self.layer[1]( - hidden_states=hidden_states, attn_metadata=attn_metadata - ) + hidden_states = self.layer[1](hidden_states=hidden_states, attn_metadata=attn_metadata) # Apply Feed Forward layer hidden_states = self.layer[2](hidden_states) @@ -477,6 +428,7 @@ class T5Block(nn.Module): class T5Stack(nn.Module): + def __init__( self, config: T5Config, @@ -491,35 +443,27 @@ class T5Stack(nn.Module): self.embed_tokens = embed_tokens self.is_umt5 = is_umt5 if is_umt5: - self.block = nn.ModuleList( - [ - T5Block( - config, - is_decoder=is_decoder, - has_relative_attention_bias=True, - quant_config=quant_config, - prefix=f"{prefix}.blocks.{i}", - ) - for i in range(n_layers) - ] - ) + self.block = nn.ModuleList([ + T5Block( + config, + is_decoder=is_decoder, + has_relative_attention_bias=True, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{i}", + ) for i in range(n_layers) + ]) else: # Only the first block has relative positional encoding. - self.block = nn.ModuleList( - [ - T5Block( - config, - is_decoder=is_decoder, - has_relative_attention_bias=i == 0, - quant_config=quant_config, - prefix=f"{prefix}.blocks.{i}", - ) - for i in range(n_layers) - ] - ) - self.final_layer_norm = RMSNorm( - config.d_model, eps=config.layer_norm_epsilon - ) + self.block = nn.ModuleList([ + T5Block( + config, + is_decoder=is_decoder, + has_relative_attention_bias=i == 0, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{i}", + ) for i in range(n_layers) + ]) + self.final_layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon) def forward( self, @@ -540,6 +484,7 @@ class T5Stack(nn.Module): class T5EncoderModel(TextEncoder): + def __init__(self, config: T5Config, prefix: str = ""): super().__init__(config) @@ -583,9 +528,7 @@ class T5EncoderModel(TextEncoder): attention_mask=attention_mask, ) - def load_weights( - self, weights: Iterable[tuple[str, torch.Tensor]] - ) -> set[str]: + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ # (param_name, shard_name, shard_id) (".qkv_proj", ".q", "q"), @@ -623,15 +566,14 @@ class T5EncoderModel(TextEncoder): continue param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) + weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params class UMT5EncoderModel(TextEncoder): + def __init__(self, config: T5Config, prefix: str = ""): super().__init__(config) @@ -675,9 +617,7 @@ class UMT5EncoderModel(TextEncoder): attention_mask=attention_mask, ) - def load_weights( - self, weights: Iterable[tuple[str, torch.Tensor]] - ) -> set[str]: + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: @@ -685,9 +625,9 @@ class UMT5EncoderModel(TextEncoder): if "decoder" in name or "lm_head" in name: continue for ( - param_name, - weight_name, - shard_id, + param_name, + weight_name, + shard_id, ) in self.config.arch_config.stacked_params_mapping: if weight_name not in name: continue @@ -713,12 +653,11 @@ class UMT5EncoderModel(TextEncoder): continue param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) + weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params + # Entry point for model registry EntryClass = [UMT5EncoderModel, T5EncoderModel] diff --git a/fastvideo/models/encoders/t5_hf.py b/fastvideo/models/encoders/t5_hf.py index 8e525e5b2..920399c89 100644 --- a/fastvideo/models/encoders/t5_hf.py +++ b/fastvideo/models/encoders/t5_hf.py @@ -11,7 +11,7 @@ from fastvideo.models.encoders.base import TextEncoder class T5EncoderModel(TextEncoder): - + supports_hf_from_pretrained: bool = True def __init__(self, config: T5Config, hf_model: Any | None = None) -> None: @@ -37,9 +37,7 @@ class T5EncoderModel(TextEncoder): hf = HFT5EncoderModel.from_pretrained(model_path, dtype=dtype, **kwargs) except TypeError: # Backward-compatible fallback for older Transformers versions. - hf = HFT5EncoderModel.from_pretrained( - model_path, torch_dtype=dtype, **kwargs - ) + hf = HFT5EncoderModel.from_pretrained(model_path, torch_dtype=dtype, **kwargs) hf = hf.eval().to(device=device, dtype=dtype) return cls(config=config, hf_model=hf).eval() @@ -54,18 +52,14 @@ class T5EncoderModel(TextEncoder): **kwargs, ) -> BaseEncoderOutput: if self.hf_model is None: - raise RuntimeError( - "T5EncoderModel(HF) is not initialized. Use " - "`from_pretrained_local(...)` to construct a loaded instance." - ) + raise RuntimeError("T5EncoderModel(HF) is not initialized. Use " + "`from_pretrained_local(...)` to construct a loaded instance.") out = self.hf_model( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_hidden_states=bool(output_hidden_states) - if output_hidden_states is not None - else False, + output_hidden_states=bool(output_hidden_states) if output_hidden_states is not None else False, return_dict=True, ) return BaseEncoderOutput( diff --git a/fastvideo/models/encoders/t5gemma.py b/fastvideo/models/encoders/t5gemma.py index 8977e3230..a67088e80 100644 --- a/fastvideo/models/encoders/t5gemma.py +++ b/fastvideo/models/encoders/t5gemma.py @@ -61,11 +61,9 @@ class T5GemmaEncoderModel(TextEncoder): path = self.t5gemma_model_path if not path: - raise ValueError( - "t5gemma_model_path must be set. Expected " - "`google/t5gemma-9b-9b-ul2` or a local path to an " - "equivalent T5-Gemma encoder." - ) + raise ValueError("t5gemma_model_path must be set. Expected " + "`google/t5gemma-9b-9b-ul2` or a local path to an " + "equivalent T5-Gemma encoder.") dtype = getattr(torch, self.t5gemma_dtype, torch.bfloat16) model = HFEncoder.from_pretrained( path, @@ -124,9 +122,7 @@ class T5GemmaEncoderModel(TextEncoder): attention_mask=attention_mask, ) - def load_weights( - self, weights: Iterable[tuple[str, torch.Tensor]] - ) -> set[str]: + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # The HF T5-Gemma encoder is lazy-loaded from `t5gemma_model_path` # (see `_build_t5gemma_model`), so this wrapper owns zero # FastVideo-native parameters and `named_parameters()` is filtered diff --git a/fastvideo/models/encoders/vision.py b/fastvideo/models/encoders/vision.py index de3ab5739..603033262 100644 --- a/fastvideo/models/encoders/vision.py +++ b/fastvideo/models/encoders/vision.py @@ -79,8 +79,7 @@ def resolve_visual_encoder_outputs( num_loaded_layers = len(encoder_outputs) - 1 offset = max_possible_layers - num_loaded_layers hs_pool = [ - encoder_outputs[layer_idx] - if layer_idx >= 0 else encoder_outputs[layer_idx + offset] + encoder_outputs[layer_idx] if layer_idx >= 0 else encoder_outputs[layer_idx + offset] for layer_idx in feature_sample_layers ] @@ -88,4 +87,4 @@ def resolve_visual_encoder_outputs( uses_last_layer = feature_sample_layers[-1] in (len(hs_pool) - 1, -1) if post_layer_norm is not None and uses_last_layer: hs_pool[-1] = post_layer_norm(encoder_outputs) - return torch.cat(hs_pool, dim=-1) \ No newline at end of file + return torch.cat(hs_pool, dim=-1) diff --git a/fastvideo/models/hf_transformer_utils.py b/fastvideo/models/hf_transformer_utils.py index 8e124a95f..5f26179d9 100644 --- a/fastvideo/models/hf_transformer_utils.py +++ b/fastvideo/models/hf_transformer_utils.py @@ -24,8 +24,7 @@ from typing import Any from huggingface_hub import snapshot_download from transformers import AutoConfig, PretrainedConfig -from transformers.models.auto.modeling_auto import ( - MODEL_FOR_CAUSAL_LM_MAPPING_NAMES) +from transformers.models.auto.modeling_auto import (MODEL_FOR_CAUSAL_LM_MAPPING_NAMES) _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = { # ChatGLMConfig.model_type: ChatGLMConfig, @@ -43,8 +42,7 @@ def download_from_hf(model_path: str): if os.path.exists(model_path): return model_path - return snapshot_download(model_path, - allow_patterns=["*.json", "*.bin", "*.model"]) + return snapshot_download(model_path, allow_patterns=["*.json", "*.bin", "*.model"]) def get_hf_config( @@ -58,10 +56,7 @@ def get_hf_config( if is_gguf: raise NotImplementedError("GGUF models are not supported.") - config = AutoConfig.from_pretrained(model, - trust_remote_code=trust_remote_code, - revision=revision, - **kwargs) + config = AutoConfig.from_pretrained(model, trust_remote_code=trust_remote_code, revision=revision, **kwargs) if config.model_type in _CONFIG_REGISTRY: config_class = _CONFIG_REGISTRY[config.model_type] config = config_class.from_pretrained(model, revision=revision) @@ -73,8 +68,7 @@ def get_hf_config( # Special architecture mapping check for GGUF models if is_gguf: if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES: - raise RuntimeError( - f"Can't get gguf config for {config.model_type}.") + raise RuntimeError(f"Can't get gguf config for {config.model_type}.") model_type = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type] config.update({"architectures": [model_type]}) @@ -107,9 +101,7 @@ def get_diffusers_config(model: str, ) -> dict[str, Any]: # TODO(will): apply any overrides from inference args return config_dict except Exception as e: - raise RuntimeError( - f"Failed to load diffusers config from {config_file}: {e}" - ) from e + raise RuntimeError(f"Failed to load diffusers config from {config_file}: {e}") from e raise RuntimeError(f"Config file not found at {config_file}") else: raise RuntimeError(f"Diffusers config file not found at {model}") @@ -131,8 +123,7 @@ CONTEXT_LENGTH_KEYS = [ def attach_additional_stop_token_ids(tokenizer): # Special handling for stop token <|eom_id|> generated by llama 3 tool use. if "<|eom_id|>" in tokenizer.get_added_vocab(): - tokenizer.additional_stop_token_ids = set( - [tokenizer.get_added_vocab()["<|eom_id|>"]]) + tokenizer.additional_stop_token_ids = set([tokenizer.get_added_vocab()["<|eom_id|>"]]) else: tokenizer.additional_stop_token_ids = None diff --git a/fastvideo/models/loader/benchmarks/benchmark_weight_loading.py b/fastvideo/models/loader/benchmarks/benchmark_weight_loading.py index 360ab3542..0ce6d64b8 100644 --- a/fastvideo/models/loader/benchmarks/benchmark_weight_loading.py +++ b/fastvideo/models/loader/benchmarks/benchmark_weight_loading.py @@ -50,19 +50,16 @@ def benchmark_loading( # Count total params and bytes on first pass total_params = 0 total_bytes = 0 - for name, tensor in safetensors_weights_iterator( - files, to_cpu=to_cpu, broadcast=broadcast): + for name, tensor in safetensors_weights_iterator(files, to_cpu=to_cpu, broadcast=broadcast): total_params += 1 total_bytes += tensor.nelement() * tensor.element_size() if rank == 0: - logger.info("[%s] %d tensors, %.2f GB total", - label, total_params, total_bytes / 1e9) + logger.info("[%s] %d tensors, %.2f GB total", label, total_params, total_bytes / 1e9) # Warmup for _ in range(warmup): - for _ in safetensors_weights_iterator( - files, to_cpu=to_cpu, broadcast=broadcast): + for _ in safetensors_weights_iterator(files, to_cpu=to_cpu, broadcast=broadcast): pass if dist.is_initialized(): dist.barrier() @@ -75,8 +72,7 @@ def benchmark_loading( torch.cuda.synchronize() if torch.cuda.is_available() else None t0 = time.perf_counter() - for _ in safetensors_weights_iterator( - files, to_cpu=to_cpu, broadcast=broadcast): + for _ in safetensors_weights_iterator(files, to_cpu=to_cpu, broadcast=broadcast): pass torch.cuda.synchronize() if torch.cuda.is_available() else None @@ -92,27 +88,30 @@ def benchmark_loading( logger.info( "[%s] avg %.3fs | best %.3fs | throughput %.2f GB/s " "(over %d runs, %d warmup, %d GPU(s), node_size=%d)", - label, avg, best, throughput, - repeats, warmup, node_group.world_size, node_group.world_size, + label, + avg, + best, + throughput, + repeats, + warmup, + node_group.world_size, + node_group.world_size, ) def main() -> None: - parser = argparse.ArgumentParser( - description="Benchmark model weight loading speed") - parser.add_argument( - "--model-path", type=str, required=True, - help="Path to a local model directory containing .safetensors files") - parser.add_argument( - "--subfolder", type=str, default=None, - help="Subfolder within model-path to load (e.g. 'transformer', " - "'text_encoder'). If not set, loads from model-path directly.") - parser.add_argument( - "--warmup", type=int, default=1, - help="Number of warmup iterations (default: 1)") - parser.add_argument( - "--repeats", type=int, default=3, - help="Number of timed iterations (default: 3)") + parser = argparse.ArgumentParser(description="Benchmark model weight loading speed") + parser.add_argument("--model-path", + type=str, + required=True, + help="Path to a local model directory containing .safetensors files") + parser.add_argument("--subfolder", + type=str, + default=None, + help="Subfolder within model-path to load (e.g. 'transformer', " + "'text_encoder'). If not set, loads from model-path directly.") + parser.add_argument("--warmup", type=int, default=1, help="Number of warmup iterations (default: 1)") + parser.add_argument("--repeats", type=int, default=3, help="Number of timed iterations (default: 3)") args = parser.parse_args() load_path = args.model_path @@ -141,13 +140,21 @@ def main() -> None: logger.info("World size: %d", world_size) # Benchmark: load to CPU (no broadcast, each rank reads independently) - benchmark_loading(files, to_cpu=True, broadcast=False, warmup=args.warmup, - repeats=args.repeats, label="to_cpu=True") + benchmark_loading(files, + to_cpu=True, + broadcast=False, + warmup=args.warmup, + repeats=args.repeats, + label="to_cpu=True") # Benchmark: load to GPU (rank 0 reads, broadcasts to others) if torch.cuda.is_available(): - benchmark_loading(files, to_cpu=False, broadcast=True, warmup=args.warmup, - repeats=args.repeats, label="to_cpu=False (broadcast)") + benchmark_loading(files, + to_cpu=False, + broadcast=True, + warmup=args.warmup, + repeats=args.repeats, + label="to_cpu=False (broadcast)") cleanup_dist_env_and_memory() diff --git a/fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py b/fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py index f6c4a0926..974f74024 100644 --- a/fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py +++ b/fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py @@ -47,8 +47,7 @@ def load_independent(files: list[str], device: str): yield name, param -def load_broadcast(files: list[str], device: str, node_group, - async_op: bool = False): +def load_broadcast(files: list[str], device: str, node_group, async_op: bool = False): """After-PR behavior: rank 0 reads from disk, broadcasts to other ranks.""" local_rank = node_group.local_rank handles = [] @@ -118,19 +117,16 @@ def run_benchmark(label, loader_fn, files, warmup, repeats): avg = sum(times) / len(times) best = min(times) throughput = total_bytes / avg / 1e9 - print( - f" {label:30s} " - f"avg {avg:.3f}s best {best:.3f}s " - f"throughput {throughput:.2f} GB/s " - f"({total_params} tensors, {total_bytes/1e9:.2f} GB)" - ) + print(f" {label:30s} " + f"avg {avg:.3f}s best {best:.3f}s " + f"throughput {throughput:.2f} GB/s " + f"({total_params} tensors, {total_bytes/1e9:.2f} GB)") return avg, best, throughput, total_bytes return None, None, None, None def main(): - parser = argparse.ArgumentParser( - description="A/B benchmark: independent vs broadcast weight loading") + parser = argparse.ArgumentParser(description="A/B benchmark: independent vs broadcast weight loading") parser.add_argument("--model-path", type=str, required=True) parser.add_argument("--subfolder", type=str, default=None) parser.add_argument("--warmup", type=int, default=1) @@ -161,7 +157,9 @@ def main(): before_avg, before_best, before_tp, total_bytes = run_benchmark( "before (independent read)", lambda f: load_independent(f, device), - files, args.warmup, args.repeats, + files, + args.warmup, + args.repeats, ) # Clear GPU memory between benchmarks @@ -172,7 +170,9 @@ def main(): after_sync_avg, after_sync_best, after_sync_tp, _ = run_benchmark( "after (sync broadcast)", lambda f: load_broadcast(f, device, node_group, async_op=False), - files, args.warmup, args.repeats, + files, + args.warmup, + args.repeats, ) # Clear GPU memory between benchmarks @@ -183,7 +183,9 @@ def main(): after_async_avg, after_async_best, after_async_tp, _ = run_benchmark( "after (async broadcast)", lambda f: load_broadcast(f, device, node_group, async_op=True), - files, args.warmup, args.repeats, + files, + args.warmup, + args.repeats, ) if rank == 0: diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index dba975819..9075de621 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -68,9 +68,7 @@ class ComponentLoader(ABC): raise NotImplementedError @classmethod - def for_module_type( - cls, module_type: str, transformers_or_diffusers: str - ) -> "ComponentLoader": + def for_module_type(cls, module_type: str, transformers_or_diffusers: str) -> "ComponentLoader": """ Factory method to create a component loader for a specific module type. @@ -196,24 +194,16 @@ class TextEncoderLoader(ComponentLoader): break if use_safetensors: - hf_weights_files = filter_duplicate_safetensors_files( - hf_weights_files, hf_folder, index_file - ) + hf_weights_files = filter_duplicate_safetensors_files(hf_weights_files, hf_folder, index_file) else: - hf_weights_files = filter_files_not_needed_for_inference( - hf_weights_files - ) + hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) if len(hf_weights_files) == 0: - raise RuntimeError( - f"Cannot find any model weights with `{model_name_or_path}`" - ) + raise RuntimeError(f"Cannot find any model weights with `{model_name_or_path}`") return hf_folder, hf_weights_files, use_safetensors - def _get_weights_iterator( - self, source: "Source", to_cpu: bool - ) -> Generator[tuple[str, torch.Tensor], None, None]: + def _get_weights_iterator(self, source: "Source", to_cpu: bool) -> Generator[tuple[str, torch.Tensor], None, None]: """Get an iterator for the model weights based on the load format.""" hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( source.model_or_path, @@ -221,21 +211,14 @@ class TextEncoderLoader(ComponentLoader): source.allow_patterns_overrides, ) if use_safetensors: - weights_iterator = safetensors_weights_iterator( - hf_weights_files, to_cpu=to_cpu - ) + weights_iterator = safetensors_weights_iterator(hf_weights_files, to_cpu=to_cpu) else: - weights_iterator = pt_weights_iterator( - hf_weights_files, to_cpu=to_cpu - ) + weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu=to_cpu) if self.counter_before_loading_weights == 0.0: self.counter_before_loading_weights = time.perf_counter() # Apply the prefix. - return ( - (source.prefix + name, tensor) - for (name, tensor) in weights_iterator - ) + return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) def _get_all_weights( self, @@ -247,9 +230,7 @@ class TextEncoderLoader(ComponentLoader): model_path, prefix="", fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True), - allow_patterns_overrides=getattr( - model, "allow_patterns_overrides", None - ), + allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None), ) yield from self._get_weights_iterator(primary_weights, to_cpu) @@ -293,20 +274,14 @@ class TextEncoderLoader(ComponentLoader): model_config["gemma_model_path"] = gemma_path if gemma_path and not gemma_path_from_candidate: if not os.path.isabs(gemma_path): - model_config["gemma_model_path"] = os.path.normpath( - os.path.join(repo_root, gemma_path) - ) - transformer_config_path = os.path.join( - repo_root, "transformer", "config.json" - ) + model_config["gemma_model_path"] = os.path.normpath(os.path.join(repo_root, gemma_path)) + transformer_config_path = os.path.join(repo_root, "transformer", "config.json") if os.path.isfile(transformer_config_path): try: with open(transformer_config_path, encoding="utf-8") as f: transformer_config = json.load(f) - if ( - "connector_double_precision_rope" not in model_config - or not model_config["connector_double_precision_rope"] - ): + if ("connector_double_precision_rope" not in model_config + or not model_config["connector_double_precision_rope"]): if transformer_config.get("double_precision_rope") is True: model_config["connector_double_precision_rope"] = True if "connector_rope_type" not in model_config: @@ -331,17 +306,11 @@ class TextEncoderLoader(ComponentLoader): f"text encoder index {idx} out of range for text_encoder_configs (len={len(encoder_configs)}), model_path={model_path}" ) encoder_config = encoder_configs[idx] - if ( - model_config.get("architectures") == ["CLIPModel"] - and isinstance(model_config.get("text_config"), dict) - ): - valid_arch_fields = { - f.name for f in dataclasses.fields(encoder_config.arch_config) - } + if (model_config.get("architectures") == ["CLIPModel"] and isinstance(model_config.get("text_config"), dict)): + valid_arch_fields = {f.name for f in dataclasses.fields(encoder_config.arch_config)} model_config = { key: value - for key, value in deepcopy(model_config["text_config"]).items() - if key in valid_arch_fields + for key, value in deepcopy(model_config["text_config"]).items() if key in valid_arch_fields } model_config["architectures"] = ["CLIPTextModel"] encoder_config.update_model_arch(model_config) @@ -364,40 +333,27 @@ class TextEncoderLoader(ComponentLoader): ) def load_model( - self, - model_path: str, - model_config: EncoderConfig, - target_device: torch.device, - fastvideo_args: FastVideoArgs, - dtype: str = "fp16", - use_text_encoder_override: bool = False, # prevent subclasses from misusing + self, + model_path: str, + model_config: EncoderConfig, + target_device: torch.device, + fastvideo_args: FastVideoArgs, + dtype: str = "fp16", + use_text_encoder_override: bool = False, # prevent subclasses from misusing ): - use_cpu_offload = ( - fastvideo_args.text_encoder_cpu_offload - and len(getattr(model_config, "_fsdp_shard_conditions", [])) > 0 - ) + use_cpu_offload = (fastvideo_args.text_encoder_cpu_offload + and len(getattr(model_config, "_fsdp_shard_conditions", [])) > 0) from fastvideo.platforms import current_platform if fastvideo_args.text_encoder_cpu_offload: - target_device = ( - torch.device("mps") - if current_platform.is_mps() - else torch.device("cpu") - ) + target_device = (torch.device("mps") if current_platform.is_mps() else torch.device("cpu")) # Set quantization config if specified - if ( - use_text_encoder_override - and fastvideo_args.override_text_encoder_quant is not None - ): + if (use_text_encoder_override and fastvideo_args.override_text_encoder_quant is not None): if fastvideo_args.override_text_encoder_safetensors is None: - raise ValueError( - "override_text_encoder_quant is set but override_text_encoder_safetensors is None" - ) - quant_cls = get_quantization_config( - fastvideo_args.override_text_encoder_quant - ) + raise ValueError("override_text_encoder_quant is set but override_text_encoder_safetensors is None") + quant_cls = get_quantization_config(fastvideo_args.override_text_encoder_quant) model_config.quant_config = quant_cls() with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]): @@ -420,30 +376,24 @@ class TextEncoderLoader(ComponentLoader): model = model_cls(model_config) # type: ignore weights_to_load = {name for name, _ in model.named_parameters()} - if ( - use_text_encoder_override - and fastvideo_args.override_text_encoder_safetensors is not None - ): + if (use_text_encoder_override and fastvideo_args.override_text_encoder_safetensors is not None): loaded_weights: set[str] = model.load_weights( safetensors_weights_iterator( [fastvideo_args.override_text_encoder_safetensors], to_cpu=use_cpu_offload, - ) - ) # type: ignore + )) # type: ignore else: loaded_weights: set[str] = model.load_weights( self._get_all_weights( model, model_path, to_cpu=fastvideo_args.text_encoder_cpu_offload, - ) - ) # type: ignore + )) # type: ignore self.counter_after_loading_weights = time.perf_counter() logger.info( "Loading weights took %.2f seconds", - self.counter_after_loading_weights - - self.counter_before_loading_weights, + self.counter_after_loading_weights - self.counter_before_loading_weights, ) # Explicitly move model to target device after loading weights @@ -455,9 +405,7 @@ class TextEncoderLoader(ComponentLoader): pin_cpu_memory = fastvideo_args.pin_cpu_memory and is_pin_memory_available() # Disable FSDP for MPS as it's not compatible if current_platform.is_mps(): - logger.info( - "Disabling FSDP sharding for MPS platform as it's not compatible" - ) + logger.info("Disabling FSDP sharding for MPS platform as it's not compatible") elif current_platform.is_npu(): mesh = init_device_mesh( "npu", @@ -491,15 +439,14 @@ class TextEncoderLoader(ComponentLoader): # if loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights if weights_not_loaded and model_config.quant_config is None: - raise ValueError( - "Following weights were not initialized from " - f"checkpoint: {weights_not_loaded}" - ) + raise ValueError("Following weights were not initialized from " + f"checkpoint: {weights_not_loaded}") return model.eval() class ImageEncoderLoader(TextEncoderLoader): + def load(self, model_path: str, fastvideo_args: FastVideoArgs): """Load the text encoders based on the model path, and inference args.""" # model_config: PretrainedConfig = get_hf_config( @@ -523,11 +470,7 @@ class ImageEncoderLoader(TextEncoderLoader): from fastvideo.platforms import current_platform if fastvideo_args.image_encoder_cpu_offload: - target_device = ( - torch.device("mps") - if current_platform.is_mps() - else torch.device("cpu") - ) + target_device = (torch.device("mps") if current_platform.is_mps() else torch.device("cpu")) else: target_device = get_local_torch_device() # TODO(will): add support for other dtypes @@ -545,8 +488,7 @@ class VisionLanguageEncoderLoader(ComponentLoader): def load(self, model_path: str, fastvideo_args: FastVideoArgs): from fastvideo.distributed.parallel_state import get_local_torch_device - from fastvideo.models.encoders.glm_image_ar_loader import ( - GlmImageARLoader) + from fastvideo.models.encoders.glm_image_ar_loader import (GlmImageARLoader) logger.info("Loading vision-language encoder from %s", model_path) target_device = get_local_torch_device() @@ -580,12 +522,8 @@ class ImageProcessorLoader(ComponentLoader): """Load the image processor based on the model path, and inference args.""" logger.info("Loading image processor from %s", model_path) - image_processor = AutoImageProcessor.from_pretrained( - model_path, - ) - logger.info( - "Loaded image processor: %s", image_processor.__class__.__name__ - ) + image_processor = AutoImageProcessor.from_pretrained(model_path, ) + logger.info("Loaded image processor: %s", image_processor.__class__.__name__) return image_processor @@ -600,9 +538,7 @@ class TokenizerLoader(ComponentLoader): # LTX2 checkpoints may not ship a top-level tokenizer/ directory. # In that case, tokenizer assets live under text_encoder/gemma/. if not os.path.isdir(resolved_model_path): - ltx2_gemma_path = os.path.normpath( - os.path.join(resolved_model_path, "..", "text_encoder", - "gemma")) + ltx2_gemma_path = os.path.normpath(os.path.join(resolved_model_path, "..", "text_encoder", "gemma")) if os.path.isdir(ltx2_gemma_path): logger.info( "Tokenizer directory %s missing; falling back to %s", @@ -618,10 +554,8 @@ class TokenizerLoader(ComponentLoader): try: with open(tokenizer_cfg_path, "r") as f: tokenizer_cfg = json.load(f) - if isinstance(tokenizer_cfg, dict) and ( - tokenizer_cfg.get("_class_name") == "AutoProcessor" - or "processor_type" in tokenizer_cfg - ): + if isinstance(tokenizer_cfg, dict) and (tokenizer_cfg.get("_class_name") == "AutoProcessor" + or "processor_type" in tokenizer_cfg): src = tokenizer_cfg.get("_name_or_path", "") if isinstance(src, str) and src.strip(): processor = AutoProcessor.from_pretrained( @@ -688,9 +622,7 @@ class TokenizerLoader(ComponentLoader): padding_side = None if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"): try: - arch_config = fastvideo_args.pipeline_config.text_encoder_configs[ - 0 - ].arch_config + arch_config = fastvideo_args.pipeline_config.text_encoder_configs[0].arch_config padding_side = getattr(arch_config, "padding_side", None) except Exception: padding_side = None @@ -710,8 +642,7 @@ class VAELoader(ComponentLoader): """Locate tokenizer-backed VAE checkpoint used by GEN3C integration.""" candidates = [ os.path.join(model_path, "tokenizer.pth"), - os.path.join(os.path.dirname(model_path), "tokenizer", - "tokenizer.pth"), + os.path.join(os.path.dirname(model_path), "tokenizer", "tokenizer.pth"), ] for candidate in candidates: if os.path.exists(candidate): @@ -727,8 +658,7 @@ class VAELoader(ComponentLoader): ] required = ("encoder.jit", "decoder.jit", "mean_std.pt") for directory in candidates: - if all(os.path.exists(os.path.join(directory, name)) - for name in required): + if all(os.path.exists(os.path.join(directory, name)) for name in required): return directory return None @@ -738,46 +668,32 @@ class VAELoader(ComponentLoader): class_name = config.pop("_class_name") config.pop("_name_or_path", None) assert class_name is not None, ( - "Model config does not contain a _class_name attribute. Only diffusers format is supported." - ) + "Model config does not contain a _class_name attribute. Only diffusers format is supported.") fastvideo_args.model_paths["vae"] = model_path from fastvideo.platforms import current_platform if fastvideo_args.vae_cpu_offload: - target_device = ( - torch.device("mps") - if current_platform.is_mps() - else torch.device("cpu") - ) + target_device = (torch.device("mps") if current_platform.is_mps() else torch.device("cpu")) else: target_device = get_local_torch_device() - with set_default_torch_dtype( - PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] - if fastvideo_args.pipeline_config.vae_precision - else torch.bfloat16 - ): + with set_default_torch_dtype(PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] if fastvideo_args. + pipeline_config.vae_precision else torch.bfloat16): pipeline_name = fastvideo_args.pipeline_config.__class__.__name__ is_gen3c = pipeline_name.startswith("Gen3C") is_cosmos25 = pipeline_name == "Cosmos25Config" # GEN3C: prefer tokenizer-backed VAE checkpoint when available. # This aligns latent conditioning with the GEN3C temporal contract. - if is_gen3c and class_name in ( - "AutoencoderKLWan", "AutoencoderKLGen3CTokenizer"): - from fastvideo.models.vaes.gen3c_tokenizer_vae import ( - AutoencoderKLGen3CTokenizer) + if is_gen3c and class_name in ("AutoencoderKLWan", "AutoencoderKLGen3CTokenizer"): + from fastvideo.models.vaes.gen3c_tokenizer_vae import (AutoencoderKLGen3CTokenizer) - dtype = PRECISION_TO_TYPE[ - fastvideo_args.pipeline_config.vae_precision] - num_frames = int( - getattr(fastvideo_args.pipeline_config, "num_frames", 121)) - state_t = int( - getattr(fastvideo_args.pipeline_config, "state_t", 16)) + dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] + num_frames = int(getattr(fastvideo_args.pipeline_config, "num_frames", 121)) + state_t = int(getattr(fastvideo_args.pipeline_config, "state_t", 16)) if state_t > 1 and num_frames > 1: - target_temporal = max(1, - (num_frames - 1) // (state_t - 1)) + target_temporal = max(1, (num_frames - 1) // (state_t - 1)) else: target_temporal = 8 @@ -797,8 +713,7 @@ class VAELoader(ComponentLoader): ) return vae.eval() - tokenizer_ckpt = self._find_gen3c_tokenizer_checkpoint( - model_path) + tokenizer_ckpt = self._find_gen3c_tokenizer_checkpoint(model_path) if tokenizer_ckpt is not None: vae = AutoencoderKLGen3CTokenizer.from_tokenizer_checkpoint( tokenizer_ckpt, @@ -828,9 +743,7 @@ class VAELoader(ComponentLoader): weight_path = os.path.join(model_path, "tokenizer.safetensors") if not os.path.exists(weight_path): - raise FileNotFoundError( - f"Missing Cosmos2.5 VAE weights: {weight_path}" - ) + raise FileNotFoundError(f"Missing Cosmos2.5 VAE weights: {weight_path}") sd = safetensors_load_file(weight_path) vae.load_state_dict(sd, strict=False) return vae.eval() @@ -843,9 +756,7 @@ class VAELoader(ComponentLoader): vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) weight_path = os.path.join(model_path, "Wan2.1_VAE.pth") if not os.path.exists(weight_path): - raise FileNotFoundError( - f"Missing LingBot World 2 VAE weights: {weight_path}" - ) + raise FileNotFoundError(f"Missing LingBot World 2 VAE weights: {weight_path}") vae = vae_cls( vae_config, checkpoint_path=weight_path, @@ -860,15 +771,10 @@ class VAELoader(ComponentLoader): if hasattr(vae, "set_tiling_config"): vae_config = fastvideo_args.pipeline_config.vae_config vae.set_tiling_config( - spatial_tile_size_in_pixels=getattr( - vae_config, "ltx2_spatial_tile_size_in_pixels", 512), - spatial_tile_overlap_in_pixels=getattr( - vae_config, "ltx2_spatial_tile_overlap_in_pixels", 64), - temporal_tile_size_in_frames=getattr( - vae_config, "ltx2_temporal_tile_size_in_frames", 64), - temporal_tile_overlap_in_frames=getattr( - vae_config, - "ltx2_temporal_tile_overlap_in_frames", 24), + spatial_tile_size_in_pixels=getattr(vae_config, "ltx2_spatial_tile_size_in_pixels", 512), + spatial_tile_overlap_in_pixels=getattr(vae_config, "ltx2_spatial_tile_overlap_in_pixels", 64), + temporal_tile_size_in_frames=getattr(vae_config, "ltx2_temporal_tile_size_in_frames", 64), + temporal_tile_overlap_in_frames=getattr(vae_config, "ltx2_temporal_tile_overlap_in_frames", 24), ) else: config.pop("_class_name", None) @@ -878,8 +784,7 @@ class VAELoader(ComponentLoader): vae = vae_cls(vae_config).to(target_device) # Find all safetensors files - safetensors_list = glob.glob( - os.path.join(str(model_path), "*.safetensors")) + safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) if not safetensors_list: raise ValueError(f"No safetensors files found in {model_path}") # Common case: a single `.safetensors` checkpoint file. @@ -915,10 +820,8 @@ class VAELoader(ComponentLoader): # strictly so missing/unexpected keys are surfaced early. strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3"} vae.load_state_dict(loaded, strict=strict_load) - if (class_name == "AutoencoderKLWan" - and getattr(vae.config, "use_light_vae", False) - and target_device.type == "cuda" - and hasattr(vae, "optimize_memory_format")): + if (class_name == "AutoencoderKLWan" and getattr(vae.config, "use_light_vae", False) + and target_device.type == "cuda" and hasattr(vae, "optimize_memory_format")): vae.optimize_memory_format() return vae.eval() @@ -955,15 +858,11 @@ class AudioDecoderLoader(ComponentLoader): audio_vae.load_state_dict(loaded, strict=True) return audio_vae.eval() - precision = getattr( - fastvideo_args.pipeline_config, "audio_decoder_precision", "bf16" - ) + precision = getattr(fastvideo_args.pipeline_config, "audio_decoder_precision", "bf16") with set_default_torch_dtype(PRECISION_TO_TYPE[precision]): audio_decoder = model_cls(config).to(target_device) - safetensors_list = glob.glob( - os.path.join(str(model_path), "*.safetensors") - ) + safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) loaded: dict[str, torch.Tensor] = {} for sf_file in safetensors_list: loaded.update(safetensors_load_file(sf_file)) @@ -990,15 +889,11 @@ class VocoderLoader(ComponentLoader): model_cls, _ = ModelRegistry.resolve_model_cls(class_name) target_device = get_local_torch_device() - precision = getattr( - fastvideo_args.pipeline_config, "vocoder_precision", "bf16" - ) + precision = getattr(fastvideo_args.pipeline_config, "vocoder_precision", "bf16") with set_default_torch_dtype(PRECISION_TO_TYPE[precision]): vocoder = model_cls(config).to(target_device) - safetensors_list = glob.glob( - os.path.join(str(model_path), "*.safetensors") - ) + safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) loaded: dict[str, torch.Tensor] = {} for sf_file in safetensors_list: loaded.update(safetensors_load_file(sf_file)) @@ -1030,10 +925,8 @@ class TransformerLoader(ComponentLoader): cls_name = config.pop("_class_name") config.pop("_name_or_path", None) if cls_name is None: - raise ValueError( - "Model config does not contain a _class_name attribute. " - "Only diffusers format is supported." - ) + raise ValueError("Model config does not contain a _class_name attribute. " + "Only diffusers format is supported.") logger.info("transformer cls_name: %s", cls_name) if fastvideo_args.override_transformer_cls_name is not None: @@ -1061,9 +954,7 @@ class TransformerLoader(ComponentLoader): model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) # Find all safetensors files - safetensors_list = glob.glob( - os.path.join(str(model_path), "*.safetensors") - ) + safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) if not safetensors_list: raise ValueError(f"No safetensors files found in {model_path}") @@ -1074,31 +965,19 @@ class TransformerLoader(ComponentLoader): update_fn(weight_keys) # Check if we should use custom initialization weights - custom_weights_path = getattr( - fastvideo_args, "init_weights_from_safetensors", None - ) - use_custom_weights = ( - custom_weights_path - and os.path.exists(custom_weights_path) - and not hasattr(fastvideo_args, "_loading_teacher_critic_model") - ) + custom_weights_path = getattr(fastvideo_args, "init_weights_from_safetensors", None) + use_custom_weights = (custom_weights_path and os.path.exists(custom_weights_path) + and not hasattr(fastvideo_args, "_loading_teacher_critic_model")) if use_custom_weights: if "transformer_2" in model_path: - custom_weights_path = getattr( - fastvideo_args, "init_weights_from_safetensors_2", None - ) - assert custom_weights_path is not None, ( - "Custom initialization weights must be provided" - ) + custom_weights_path = getattr(fastvideo_args, "init_weights_from_safetensors_2", None) + assert custom_weights_path is not None, ("Custom initialization weights must be provided") if os.path.isdir(custom_weights_path): - safetensors_list = glob.glob( - os.path.join(str(custom_weights_path), "*.safetensors") - ) + safetensors_list = glob.glob(os.path.join(str(custom_weights_path), "*.safetensors")) else: assert custom_weights_path.endswith(".safetensors"), ( - "Custom initialization weights must be a safetensors file" - ) + "Custom initialization weights must be a safetensors file") safetensors_list = [custom_weights_path] logger.info( @@ -1107,22 +986,16 @@ class TransformerLoader(ComponentLoader): safetensors_list, ) - default_dtype = PRECISION_TO_TYPE[ - fastvideo_args.pipeline_config.dit_precision - ] + default_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision] # Load the model using FSDP loader - logger.info("Loading model from %s, default_dtype: %s", cls_name, - default_dtype) + logger.info("Loading model from %s, default_dtype: %s", cls_name, default_dtype) assert fastvideo_args.hsdp_shard_dim is not None # Cosmos2.5 checkpoints can include extra entries not present in the # instantiated model (e.g. pos_embedder ranges / *_extra_state). Load # non-strictly for Cosmos2.5 only; keep upstream strict behavior for others. - strict_load = not ( - cls_name.startswith("Cosmos25") - or cls_name == "Cosmos25Transformer3DModel" - or getattr(fastvideo_args.pipeline_config, "prefix", "") == "Cosmos25" - ) + strict_load = not (cls_name.startswith("Cosmos25") or cls_name == "Cosmos25Transformer3DModel" + or getattr(fastvideo_args.pipeline_config, "prefix", "") == "Cosmos25") attention_context = (_component_attention_backend_scope(None, component="transformer") if _qat_generator_only else nullcontext()) with attention_context: @@ -1133,7 +1006,10 @@ class TransformerLoader(ComponentLoader): logger.info("transformer attention backend: %s", resolved.name if resolved else "automatic selection") model = maybe_load_fsdp_model( model_cls=model_cls, - init_params={"config": dit_config, "hf_config": hf_config}, + init_params={ + "config": dit_config, + "hf_config": hf_config + }, weight_dir_list=safetensors_list, device=get_local_torch_device(), hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim, @@ -1164,17 +1040,13 @@ class TransformerLoader(ComponentLoader): if fastvideo_args.inference_mode and fastvideo_args.dit_layerwise_offload: # Check if model has nn.ModuleList for layerwise offload compatibility - has_module_list = any( - isinstance(m, nn.ModuleList) for m in model.children() - ) + has_module_list = any(isinstance(m, nn.ModuleList) for m in model.children()) if has_module_list: enable_layerwise_offload(model) else: logger.warning( "Layerwise offload requested but model %s does not have " - "nn.ModuleList structure. Skipping layerwise offload.", - cls_name - ) + "nn.ModuleList structure. Skipping layerwise offload.", cls_name) return model @@ -1187,8 +1059,7 @@ class SchedulerLoader(ComponentLoader): class_name = config.pop("_class_name") assert class_name is not None, ( - "Model config does not contain a _class_name attribute. Only diffusers format is supported." - ) + "Model config does not contain a _class_name attribute. Only diffusers format is supported.") scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name) @@ -1214,9 +1085,8 @@ class ConditionerLoader(ComponentLoader): class_name = config.pop("_class_name", None) config.pop("_name_or_path", None) if class_name is None: - raise ValueError( - f"Conditioner config at {model_path} is missing the " - f"`_class_name` attribute required to resolve a model class.") + raise ValueError(f"Conditioner config at {model_path} is missing the " + f"`_class_name` attribute required to resolve a model class.") model_cls, _ = ModelRegistry.resolve_model_cls(class_name) target_device = get_local_torch_device() @@ -1244,8 +1114,7 @@ class ConditionerLoader(ComponentLoader): weights = os.path.join(str(model_path), "diffusion_pytorch_model.safetensors") if not os.path.isfile(weights): - raise FileNotFoundError( - f"Conditioner weights not found: {weights}") + raise FileNotFoundError(f"Conditioner weights not found: {weights}") state = safetensors_load_file(weights) # Non-strict: T5 weights live outside this checkpoint (fetched in # the conditioner's `__init__` from the standard HF repo). @@ -1262,9 +1131,8 @@ class UpsamplerLoader(ComponentLoader): class_name = config_dict.pop("_class_name", None) if class_name is None: - raise ValueError( - "Model config does not contain a _class_name attribute. " - "Only diffusers format is supported.") + raise ValueError("Model config does not contain a _class_name attribute. " + "Only diffusers format is supported.") # The base PipelineConfig declares ``upsampler_config`` as a # single ``UpsamplerConfig`` instance, but Hunyuan15 narrows it @@ -1272,8 +1140,7 @@ class UpsamplerLoader(ComponentLoader): # the attribute as a multi-config when it actually is one; # otherwise the LTX-2 branch below handles the single-class # path that takes the diffusers config dict directly. - upsampler_config_attr = getattr(fastvideo_args.pipeline_config, - "upsampler_config", None) + upsampler_config_attr = getattr(fastvideo_args.pipeline_config, "upsampler_config", None) if isinstance(upsampler_config_attr, list | tuple): try: upsampler_cfg = deepcopy(upsampler_config_attr[0]) @@ -1287,21 +1154,17 @@ class UpsamplerLoader(ComponentLoader): # dict directly via LatentUpsamplerConfigurator. upsampler_cfg = deepcopy(config_dict) else: - raise AttributeError( - "pipeline_config.upsampler_config is missing; cannot build " - f"upsampler config for class {class_name}") + raise AttributeError("pipeline_config.upsampler_config is missing; cannot build " + f"upsampler config for class {class_name}") model_cls, _ = ModelRegistry.resolve_model_cls(class_name) model = model_cls(upsampler_cfg) target_device = get_local_torch_device() - upsampler_precision = getattr(fastvideo_args.pipeline_config, - "upsampler_precision", "bf16") - model = model.to(target_device, - dtype=PRECISION_TO_TYPE[upsampler_precision]) + upsampler_precision = getattr(fastvideo_args.pipeline_config, "upsampler_precision", "bf16") + model = model.to(target_device, dtype=PRECISION_TO_TYPE[upsampler_precision]) - safetensors_list = glob.glob( - os.path.join(str(model_path), "*.safetensors")) + safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) if not safetensors_list: raise ValueError(f"No safetensors files found in {model_path}") @@ -1326,6 +1189,7 @@ class UpsamplerLoader(ComponentLoader): return model.eval() + class GenericComponentLoader(ComponentLoader): """Generic loader for components that don't have a specific loader.""" @@ -1355,9 +1219,7 @@ class GenericComponentLoader(ComponentLoader): ) return model elif self.library == "diffusers": - logger.warning( - "Generic loading for diffusers components is not fully implemented" - ) + logger.warning("Generic loading for diffusers components is not fully implemented") model_config = get_diffusers_config(model=model_path) logger.info("Diffusers Model config: %s", model_config) @@ -1400,9 +1262,7 @@ class PipelineComponentLoader: ) # Get the appropriate loader for this module type - loader = ComponentLoader.for_module_type( - module_name, transformers_or_diffusers - ) + loader = ComponentLoader.for_module_type(module_name, transformers_or_diffusers) # Resolve this component's attention backend ONCE, here, from the # explicit request and the environment. A per-role request already diff --git a/fastvideo/models/loader/utils.py b/fastvideo/models/loader/utils.py index d3bbc3ce1..4783f699e 100644 --- a/fastvideo/models/loader/utils.py +++ b/fastvideo/models/loader/utils.py @@ -24,8 +24,7 @@ def set_default_torch_dtype(dtype: torch.dtype): torch.set_default_dtype(old_dtype) -def get_param_names_mapping( - mapping_dict: dict[str, str]) -> Callable[[str], tuple[str, Any, Any]]: +def get_param_names_mapping(mapping_dict: dict[str, str]) -> Callable[[str], tuple[str, Any, Any]]: """ Creates a mapping function that transforms parameter names using regex patterns. @@ -78,19 +77,13 @@ def hf_to_custom_state_dict( if isinstance(hf_param_sd, dict): hf_param_sd = hf_param_sd.items() # type: ignore for source_param_name, full_tensor in hf_param_sd: # type: ignore - target_param_name, merge_index, num_params_to_merge = param_names_mapping( - source_param_name) - reverse_param_names_mapping[target_param_name] = (source_param_name, - merge_index, - num_params_to_merge) + target_param_name, merge_index, num_params_to_merge = param_names_mapping(source_param_name) + reverse_param_names_mapping[target_param_name] = (source_param_name, merge_index, num_params_to_merge) if merge_index is not None: to_merge_params[target_param_name][merge_index] = full_tensor if len(to_merge_params[target_param_name]) == num_params_to_merge: # cat at output dim according to the merge_index order - sorted_tensors = [ - to_merge_params[target_param_name][i] - for i in range(num_params_to_merge) - ] + sorted_tensors = [to_merge_params[target_param_name][i] for i in range(num_params_to_merge)] full_tensor = torch.cat(sorted_tensors, dim=0) del to_merge_params[target_param_name] else: diff --git a/fastvideo/models/loader/weight_utils.py b/fastvideo/models/loader/weight_utils.py index 202ed4161..574a13356 100644 --- a/fastvideo/models/loader/weight_utils.py +++ b/fastvideo/models/loader/weight_utils.py @@ -82,9 +82,7 @@ def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None): # Passing both of these to the weight loader functionality breaks. # So, we use the index_file to # look up which safetensors files should be used. -def filter_duplicate_safetensors_files(hf_weights_files: list[str], - hf_folder: str, - index_file: str) -> list[str]: +def filter_duplicate_safetensors_files(hf_weights_files: list[str], hf_folder: str, index_file: str) -> list[str]: # model.safetensors.index.json is a mapping from keys in the # torch state_dict to safetensors file holding that weight. index_file_name = os.path.join(hf_folder, index_file) @@ -97,12 +95,9 @@ def filter_duplicate_safetensors_files(hf_weights_files: list[str], weight_map = json.load(f)["weight_map"] weight_files_in_index = set() for weight_name in weight_map: - weight_files_in_index.add( - os.path.join(hf_folder, weight_map[weight_name])) + weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name])) # Filter out any fields that are not found in the index file. - hf_weights_files = [ - f for f in hf_weights_files if f in weight_files_in_index - ] + hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index] return hf_weights_files @@ -111,21 +106,16 @@ SAFE_WEIGHTS_INDEX_NAME = "model.safetensors.index.json" def resolve_safetensors_files(model_path: str) -> list[str]: """Discover safetensors files in a model directory.""" - files = sorted( - glob.glob(os.path.join(model_path, "*.safetensors"))) + files = sorted(glob.glob(os.path.join(model_path, "*.safetensors"))) if not files: - raise FileNotFoundError( - f"No .safetensors files found in {model_path}") - index_file = os.path.join( - model_path, SAFE_WEIGHTS_INDEX_NAME) + raise FileNotFoundError(f"No .safetensors files found in {model_path}") + index_file = os.path.join(model_path, SAFE_WEIGHTS_INDEX_NAME) if os.path.exists(index_file): - files = filter_duplicate_safetensors_files( - files, model_path, SAFE_WEIGHTS_INDEX_NAME) + files = filter_duplicate_safetensors_files(files, model_path, SAFE_WEIGHTS_INDEX_NAME) return files -def filter_files_not_needed_for_inference( - hf_weights_files: list[str]) -> list[str]: +def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]: """ Exclude files that are not needed for inference. @@ -138,10 +128,7 @@ def filter_files_not_needed_for_inference( "scheduler.pt", "scaler.pt", ] - hf_weights_files = [ - f for f in hf_weights_files - if not any(f.endswith(x) for x in blacklist) - ] + hf_weights_files = [f for f in hf_weights_files if not any(f.endswith(x) for x in blacklist)] return hf_weights_files @@ -160,12 +147,10 @@ def _get_initialized_node_group(): return None -def safetensors_weights_iterator( - hf_weights_files: list[str], - to_cpu: bool = False, - broadcast: bool = True, - async_broadcast: bool = False -) -> Generator[tuple[str, torch.Tensor], None, None]: +def safetensors_weights_iterator(hf_weights_files: list[str], + to_cpu: bool = False, + broadcast: bool = True, + async_broadcast: bool = False) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model safetensor files. Args: hf_weights_files: List of safetensor files to load. @@ -176,8 +161,7 @@ def safetensors_weights_iterator( must iterate over all the weights before use. Only used when broadcast is True and to_cpu is False. """ node_group = _get_initialized_node_group() - local_rank = node_group.local_rank if node_group is not None else int( - os.environ.get("LOCAL_RANK", 0)) + local_rank = node_group.local_rank if node_group is not None else int(os.environ.get("LOCAL_RANK", 0)) device = str(parallel_state.get_local_torch_device()) if not to_cpu else "cpu" enable_tqdm = not torch.distributed.is_initialized() or local_rank == 0 if to_cpu or not broadcast or node_group is None: @@ -208,15 +192,12 @@ def safetensors_weights_iterator( group = node_group.device_group if async_broadcast: handle = dist.broadcast(param, - src=dist.get_global_rank( - group, 0), + src=dist.get_global_rank(group, 0), async_op=True, group=group) handles.append(handle) else: - dist.broadcast(param, - src=dist.get_global_rank(group, 0), - group=group) + dist.broadcast(param, src=dist.get_global_rank(group, 0), group=group) else: param = f.get_tensor(name) yield name, param @@ -227,11 +208,9 @@ def safetensors_weights_iterator( handles.clear() -def pt_weights_iterator( - hf_weights_files: list[str], - to_cpu: bool = False, - broadcast: bool = True -) -> Generator[tuple[str, torch.Tensor], None, None]: +def pt_weights_iterator(hf_weights_files: list[str], + to_cpu: bool = False, + broadcast: bool = True) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model bin/pt files. Args: @@ -241,8 +220,7 @@ def pt_weights_iterator( torch.load and do not use the safetensors broadcast path. """ node_group = _get_initialized_node_group() - local_rank = node_group.local_rank if node_group is not None else int( - os.environ.get("LOCAL_RANK", 0)) + local_rank = node_group.local_rank if node_group is not None else int(os.environ.get("LOCAL_RANK", 0)) device = str(parallel_state.get_local_torch_device()) if not to_cpu else "cpu" enable_tqdm = not torch.distributed.is_initialized() or local_rank == 0 for bin_file in tqdm( @@ -256,8 +234,7 @@ def pt_weights_iterator( del state -def default_weight_loader(param: torch.Tensor, - loaded_weight: torch.Tensor) -> None: +def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: """Default weight loader.""" try: if param.numel() == 1 and loaded_weight.numel() == 1: @@ -266,9 +243,8 @@ def default_weight_loader(param: torch.Tensor, # "broadcast" instead of copy param.data.fill_(loaded_weight.item()) else: - assert param.size() == loaded_weight.size(), ( - f"Attempted to load weight ({loaded_weight.size()}) " - f"into parameter ({param.size()})") + assert param.size() == loaded_weight.size(), (f"Attempted to load weight ({loaded_weight.size()}) " + f"into parameter ({param.size()})") param.data.copy_(loaded_weight) except Exception: @@ -295,42 +271,35 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None: None: If the remapped name is not found in params_dict. """ if name.endswith(".kv_scale"): - logger.warning_once( - "DEPRECATED. Found kv_scale in the checkpoint. " - "This format is deprecated in favor of separate k_scale and " - "v_scale tensors and will be removed in a future release. " - "Functionally, we will remap kv_scale to k_scale and duplicate " - "k_scale to v_scale") + logger.warning_once("DEPRECATED. Found kv_scale in the checkpoint. " + "This format is deprecated in favor of separate k_scale and " + "v_scale tensors and will be removed in a future release. " + "Functionally, we will remap kv_scale to k_scale and duplicate " + "k_scale to v_scale") # NOTE: we remap the deprecated kv_scale to k_scale remapped_name = name.replace(".kv_scale", ".attn.k_scale") if remapped_name not in params_dict: - logger.warning_once( - f"Found kv_scale in the checkpoint (e.g. {name}), " - "but not found the expected name in the model " - f"(e.g. {remapped_name}). kv_scale is " - "not loaded.") + logger.warning_once(f"Found kv_scale in the checkpoint (e.g. {name}), " + "but not found the expected name in the model " + f"(e.g. {remapped_name}). kv_scale is " + "not loaded.") return None return remapped_name possible_scale_names = [".k_scale", ".v_scale"] - modelopt_scale_names = [ - ".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale" - ] + modelopt_scale_names = [".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale"] for scale_name in possible_scale_names: if name.endswith(scale_name): - if any(mo_scale_name in name - for mo_scale_name in modelopt_scale_names): - remapped_name = name.replace( - f".self_attn.{scale_name[1]}_proj{scale_name}", - f".self_attn.attn{scale_name}") + if any(mo_scale_name in name for mo_scale_name in modelopt_scale_names): + remapped_name = name.replace(f".self_attn.{scale_name[1]}_proj{scale_name}", + f".self_attn.attn{scale_name}") else: remapped_name = name.replace(scale_name, f".attn{scale_name}") if remapped_name not in params_dict: - logger.warning_once( - f"Found {scale_name} in the checkpoint (e.g. {name}), " - "but not found the expected name in the model " - f"(e.g. {remapped_name}). {scale_name} is " - "not loaded.") + logger.warning_once(f"Found {scale_name} in the checkpoint (e.g. {name}), " + "but not found the expected name in the model " + f"(e.g. {remapped_name}). {scale_name} is " + "not loaded.") return None return remapped_name diff --git a/fastvideo/models/mask_utils.py b/fastvideo/models/mask_utils.py index ec944c8d9..8a99b8538 100644 --- a/fastvideo/models/mask_utils.py +++ b/fastvideo/models/mask_utils.py @@ -1,6 +1,7 @@ import torch from typing import Callable, Optional + def and_masks(*mask_functions: Callable) -> Callable: """Returns a mask function that is the intersection of provided mask functions""" if not all(callable(arg) for arg in mask_functions): @@ -14,12 +15,14 @@ def and_masks(*mask_functions: Callable) -> Callable: return and_mask + def causal_mask_function(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: """ This creates a basic lower-diagonal causal mask. """ return kv_idx <= q_idx + def padding_mask_function(padding_mask: torch.Tensor) -> Callable: """ This return the mask_function function corresponding to a 2D padding mask. @@ -33,9 +36,9 @@ def padding_mask_function(padding_mask: torch.Tensor) -> Callable: return inner_mask -def prepare_padding_mask( - attention_mask: Optional[torch.Tensor], kv_length: int, kv_offset: int -) -> Optional[torch.Tensor]: + +def prepare_padding_mask(attention_mask: Optional[torch.Tensor], kv_length: int, + kv_offset: int) -> Optional[torch.Tensor]: """ From the 2D attention mask, prepare the correct padding mask to use by potentially padding it. """ @@ -46,9 +49,9 @@ def prepare_padding_mask( local_padding_mask = torch.nn.functional.pad(attention_mask, (0, padding_length)) return local_padding_mask -def _non_vmap_expansion_sdpa( - batch_indices: torch.Tensor, head_indices: torch.Tensor, q_indices: torch.Tensor, kv_indices: torch.Tensor -): + +def _non_vmap_expansion_sdpa(batch_indices: torch.Tensor, head_indices: torch.Tensor, q_indices: torch.Tensor, + kv_indices: torch.Tensor): """ Used to broadcast our mask_functions over the all 4 dimensions (b_idx, h_idx, q_idx, kv_idx) of the inputs. Allows the usage of any index-based mask function without relying on vmap. @@ -64,6 +67,7 @@ def _non_vmap_expansion_sdpa( kv_indices = kv_indices[None, None, None, :] return batch_indices, head_indices, q_indices, kv_indices + def sdpa_mask( batch_size: int, cache_position: torch.Tensor, @@ -198,4 +202,4 @@ def sdpa_mask( # Expand the mask to match batch size and query length if they weren't used in the mask function attention_mask = attention_mask.expand(batch_size, -1, q_length, kv_length) - return attention_mask \ No newline at end of file + return attention_mask diff --git a/fastvideo/models/parameter.py b/fastvideo/models/parameter.py index dd7dfac10..a2c192fe2 100644 --- a/fastvideo/models/parameter.py +++ b/fastvideo/models/parameter.py @@ -60,8 +60,7 @@ class BasevLLMParameter(Parameter): return (cond1 and cond2) def _assert_and_load(self, loaded_weight: torch.Tensor) -> None: - assert (self.data.shape == loaded_weight.shape - or self._is_1d_and_scalar(loaded_weight)) + assert (self.data.shape == loaded_weight.shape or self._is_1d_and_scalar(loaded_weight)) self.data.copy_(loaded_weight) def load_column_parallel_weight(self, loaded_weight: torch.Tensor) -> None: @@ -70,8 +69,7 @@ class BasevLLMParameter(Parameter): def load_row_parallel_weight(self, loaded_weight: torch.Tensor) -> None: self._assert_and_load(loaded_weight) - def load_merged_column_weight(self, loaded_weight: torch.Tensor, - **kwargs) -> None: + def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None: self._assert_and_load(loaded_weight) def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None: @@ -100,31 +98,25 @@ class _ColumnvLLMParameter(BasevLLMParameter): def load_column_parallel_weight(self, loaded_weight: torch.Tensor) -> None: tp_rank = get_tp_rank() shard_size = self.data.shape[self.output_dim] - loaded_weight = loaded_weight.narrow(self.output_dim, - tp_rank * shard_size, shard_size) + loaded_weight = loaded_weight.narrow(self.output_dim, tp_rank * shard_size, shard_size) assert self.data.shape == loaded_weight.shape self.data.copy_(loaded_weight) - def load_merged_column_weight(self, loaded_weight: torch.Tensor, - **kwargs) -> None: + def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None: shard_offset = kwargs.get("shard_offset") shard_size = kwargs.get("shard_size") if shard_offset is None or shard_size is None: raise ValueError("shard_offset and shard_size must be provided") - if isinstance( - self, PackedColumnParameter - | PackedvLLMParameter) and self.packed_dim == self.output_dim: - shard_size, shard_offset = self.adjust_shard_indexes_for_packing( - shard_offset=shard_offset, shard_size=shard_size) + if isinstance(self, PackedColumnParameter | PackedvLLMParameter) and self.packed_dim == self.output_dim: + shard_size, shard_offset = self.adjust_shard_indexes_for_packing(shard_offset=shard_offset, + shard_size=shard_size) param_data = self.data tp_rank = get_tp_rank() - param_data = param_data.narrow(self.output_dim, shard_offset, - shard_size) - loaded_weight = loaded_weight.narrow(self.output_dim, - tp_rank * shard_size, shard_size) + param_data = param_data.narrow(self.output_dim, shard_offset, shard_size) + loaded_weight = loaded_weight.narrow(self.output_dim, tp_rank * shard_size, shard_size) assert param_data.shape == loaded_weight.shape param_data.copy_(loaded_weight) @@ -140,19 +132,15 @@ class _ColumnvLLMParameter(BasevLLMParameter): assert shard_id is not None assert num_heads is not None - if isinstance( - self, PackedColumnParameter - | PackedvLLMParameter) and self.output_dim == self.packed_dim: - shard_size, shard_offset = self.adjust_shard_indexes_for_packing( - shard_offset=shard_offset, shard_size=shard_size) + if isinstance(self, PackedColumnParameter | PackedvLLMParameter) and self.output_dim == self.packed_dim: + shard_size, shard_offset = self.adjust_shard_indexes_for_packing(shard_offset=shard_offset, + shard_size=shard_size) param_data = self.data tp_rank = get_tp_rank() shard_id = tp_rank if shard_id == "q" else tp_rank // num_heads - param_data = param_data.narrow(self.output_dim, shard_offset, - shard_size) - loaded_weight = loaded_weight.narrow(self.output_dim, - shard_id * shard_size, shard_size) + param_data = param_data.narrow(self.output_dim, shard_offset, shard_size) + loaded_weight = loaded_weight.narrow(self.output_dim, shard_id * shard_size, shard_size) assert param_data.shape == loaded_weight.shape param_data.copy_(loaded_weight) @@ -177,8 +165,7 @@ class RowvLLMParameter(BasevLLMParameter): def load_row_parallel_weight(self, loaded_weight: torch.Tensor) -> None: tp_rank = get_tp_rank() shard_size = self.data.shape[self.input_dim] - loaded_weight = loaded_weight.narrow(self.input_dim, - tp_rank * shard_size, shard_size) + loaded_weight = loaded_weight.narrow(self.input_dim, tp_rank * shard_size, shard_size) if len(loaded_weight.shape) == 0: loaded_weight = loaded_weight.reshape(1) @@ -253,8 +240,7 @@ class PerTensorScaleParameter(BasevLLMParameter): def load_column_parallel_weight(self, *args, **kwargs) -> None: super().load_row_parallel_weight(*args, **kwargs) - def _load_into_shard_id(self, loaded_weight: torch.Tensor, - shard_id: str | int, **kwargs): + def _load_into_shard_id(self, loaded_weight: torch.Tensor, shard_id: str | int, **kwargs): """ Slice the parameter data based on the shard id for loading. @@ -281,8 +267,7 @@ class PackedColumnParameter(_ColumnvLLMParameter): for more details on the packed properties. """ - def __init__(self, packed_factor: int | Fraction, packed_dim: int, - **kwargs): + def __init__(self, packed_factor: int | Fraction, packed_dim: int, **kwargs): self._packed_factor = packed_factor self._packed_dim = packed_dim super().__init__(**kwargs) @@ -295,12 +280,10 @@ class PackedColumnParameter(_ColumnvLLMParameter): def packed_factor(self): return self._packed_factor - def adjust_shard_indexes_for_packing(self, shard_size, - shard_offset) -> tuple[Any, Any]: - return _adjust_shard_indexes_for_packing( - shard_size=shard_size, - shard_offset=shard_offset, - packed_factor=self.packed_factor) + def adjust_shard_indexes_for_packing(self, shard_size, shard_offset) -> tuple[Any, Any]: + return _adjust_shard_indexes_for_packing(shard_size=shard_size, + shard_offset=shard_offset, + packed_factor=self.packed_factor) class PackedvLLMParameter(ModelWeightParameter): @@ -314,8 +297,7 @@ class PackedvLLMParameter(ModelWeightParameter): by accounting for packing and optionally, marlin tile size. """ - def __init__(self, packed_factor: int | Fraction, packed_dim: int, - **kwargs): + def __init__(self, packed_factor: int | Fraction, packed_dim: int, **kwargs): self._packed_factor = packed_factor self._packed_dim = packed_dim super().__init__(**kwargs) @@ -329,10 +311,9 @@ class PackedvLLMParameter(ModelWeightParameter): return self._packed_factor def adjust_shard_indexes_for_packing(self, shard_size, shard_offset): - return _adjust_shard_indexes_for_packing( - shard_size=shard_size, - shard_offset=shard_offset, - packed_factor=self.packed_factor) + return _adjust_shard_indexes_for_packing(shard_size=shard_size, + shard_offset=shard_offset, + packed_factor=self.packed_factor) class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): @@ -344,8 +325,7 @@ class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): pass -def permute_param_layout_(param: BasevLLMParameter, input_dim: int, - output_dim: int, **kwargs) -> BasevLLMParameter: +def permute_param_layout_(param: BasevLLMParameter, input_dim: int, output_dim: int, **kwargs) -> BasevLLMParameter: """ Permute a parameter's layout to the specified input and output dimensions, useful for forcing the parameter into a known layout, for example, if I need @@ -379,10 +359,7 @@ def permute_param_layout_(param: BasevLLMParameter, input_dim: int, # create permutation from the current layout to the layout with # self.input_dim at input_dim and self.output_dim at output_dim preserving # other dimensions - perm = [ - i for i in range(param.data.dim()) - if i not in [curr_input_dim, curr_output_dim] - ] + perm = [i for i in range(param.data.dim()) if i not in [curr_input_dim, curr_output_dim]] perm.insert(input_dim, curr_input_dim) perm.insert(output_dim, curr_output_dim) @@ -402,8 +379,7 @@ def permute_param_layout_(param: BasevLLMParameter, input_dim: int, return param -def _adjust_shard_indexes_for_packing(shard_size, shard_offset, - packed_factor) -> tuple[Any, Any]: +def _adjust_shard_indexes_for_packing(shard_size, shard_offset, packed_factor) -> tuple[Any, Any]: shard_size = shard_size // packed_factor shard_offset = shard_offset // packed_factor return shard_size, shard_offset diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index eb3308c18..00bd3a62c 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -23,21 +23,18 @@ logger = init_logger(__name__) # huggingface class name: (component_name, fastvideo module name, fastvideo class name) _TEXT_TO_VIDEO_DIT_MODELS = { - "HunyuanVideoTransformer3DModel": - ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"), - "HunyuanGameCraftTransformer3DModel": - ("dits", "hunyuangamecraft", "HunyuanGameCraftTransformer3DModel"), - "HunyuanVideo15Transformer3DModel": - ("dits", "hunyuanvideo15", "HunyuanVideo15Transformer3DModel"), - "HYWorldTransformer3DModel": - ("dits", "hyworld", "HYWorldTransformer3DModel"), + "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"), + "HunyuanGameCraftTransformer3DModel": ("dits", "hunyuangamecraft", "HunyuanGameCraftTransformer3DModel"), + "HunyuanVideo15Transformer3DModel": ("dits", "hunyuanvideo15", "HunyuanVideo15Transformer3DModel"), + "HYWorldTransformer3DModel": ("dits", "hyworld", "HYWorldTransformer3DModel"), "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"), "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"), "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"), "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"), - "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"), # Wrapper (Phase 1) + "LongCatVideoTransformer3DModel": + ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"), # Wrapper (Phase 1) "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"), # Native (Phase 2) "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"), "SD3Transformer2DModel": ("dits", "sd3", "SD3Transformer2DModel"), @@ -79,8 +76,7 @@ _TEXT_TO_IMAGE_DIT_MODELS = { _TEXT_ENCODER_MODELS = { "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"), - "CLIPTextModelWithProjection": - ("encoders", "clip", "CLIPTextModelWithProjection"), + "CLIPTextModelWithProjection": ("encoders", "clip", "CLIPTextModelWithProjection"), "LlamaModel": ("encoders", "llama", "LlamaModel"), "UMT5EncoderModel": ("encoders", "t5", "UMT5EncoderModel"), "LingBotWorld2T5EncoderModel": ("encoders", "lingbotworld2_t5", "LingBotWorld2T5EncoderModel"), @@ -88,15 +84,13 @@ _TEXT_ENCODER_MODELS = { "BertModel": ("encoders", "clip", "CLIPTextModel"), "Qwen2_5_VLTextModel": ("encoders", "qwen2_5", "Qwen2_5_VLTextModel"), "Reason1TextEncoder": ("encoders", "reason1", "Reason1TextEncoder"), - "Qwen2_5_VLForConditionalGeneration": - ("encoders", "reason1", "Reason1TextEncoder"), + "Qwen2_5_VLForConditionalGeneration": ("encoders", "reason1", "Reason1TextEncoder"), # Z-Image-Turbo's text_encoder/config.json declares architecture # "Qwen3Model"; route it to the shared Qwen3 encoder (added for Flux2 Klein). "Qwen3Model": ("encoders", "qwen3", "Qwen3ForCausalLM"), "LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"), "Qwen3ForCausalLM": ("encoders", "qwen3", "Qwen3ForCausalLM"), - "Mistral3ForConditionalGeneration": - ("encoders", "mistral3", "Mistral3ForConditionalGeneration"), + "Mistral3ForConditionalGeneration": ("encoders", "mistral3", "Mistral3ForConditionalGeneration"), } _IMAGE_ENCODER_MODELS: dict[str, tuple] = { @@ -107,16 +101,14 @@ _IMAGE_ENCODER_MODELS: dict[str, tuple] = { } _VAE_MODELS = { - "AutoencoderKLHunyuanVideo": - ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"), + "AutoencoderKLHunyuanVideo": ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"), "AutoencoderKLCausal3D": ("vaes", "gamecraftvae", "GameCraftVAE"), "AutoencoderKLHYWorld": ("vaes", "hyworldvae", "AutoencoderKLHYWorld"), "AutoencoderKLHunyuanVideo15": ("vaes", "hunyuan15vae", "AutoencoderKLHunyuanVideo15"), "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"), "LingBotWorld2WanVAE": ("vaes", "lingbotworld2_wanvae", "LingBotWorld2WanVAE"), "AutoencoderKL": ("vaes", "autoencoder_kl", "AutoencoderKL"), - "AutoencoderKLGen3CTokenizer": - ("vaes", "gen3c_tokenizer_vae", "AutoencoderKLGen3CTokenizer"), + "AutoencoderKLGen3CTokenizer": ("vaes", "gen3c_tokenizer_vae", "AutoencoderKLGen3CTokenizer"), "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo"), "CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"), "AutoencoderKLFlux2": ("vaes", "flux2vae", "AutoencoderKLFlux2"), @@ -133,17 +125,12 @@ _AUDIO_MODELS = { _SCHEDULERS = { "FlowMatchEulerDiscreteScheduler": - ("schedulers", "scheduling_flow_match_euler_discrete", - "FlowMatchEulerDiscreteScheduler"), - "UniPCMultistepScheduler": - ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"), - "FlowUniPCMultistepScheduler": - ("schedulers", "scheduling_flow_unipc_multistep", "FlowUniPCMultistepScheduler"), + ("schedulers", "scheduling_flow_match_euler_discrete", "FlowMatchEulerDiscreteScheduler"), + "UniPCMultistepScheduler": ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"), + "FlowUniPCMultistepScheduler": ("schedulers", "scheduling_flow_unipc_multistep", "FlowUniPCMultistepScheduler"), "SelfForcingFlowMatchScheduler": - ("schedulers", "scheduling_self_forcing_flow_match", - "SelfForcingFlowMatchScheduler"), - "RCMScheduler": - ("schedulers", "scheduling_rcm", "RCMScheduler"), + ("schedulers", "scheduling_self_forcing_flow_match", "SelfForcingFlowMatchScheduler"), + "RCMScheduler": ("schedulers", "scheduling_rcm", "RCMScheduler"), } _UPSAMPLERS = { @@ -171,10 +158,7 @@ MODELS_PATH = os.path.dirname(__file__) def _discover_and_register_models() -> dict[str, tuple[str, str, str]]: discovered_models: dict[str, tuple[str, str, str]] = {} for root, dirs, files in os.walk(MODELS_PATH): - dirs[:] = [ - d for d in dirs - if not d.startswith(".") and d != "__pycache__" - ] + dirs[:] = [d for d in dirs if not d.startswith(".") and d != "__pycache__"] for filename in files: if not filename.endswith(".py"): @@ -208,8 +192,7 @@ def _discover_and_register_models() -> dict[str, tuple[str, str, str]]: model_cls_name_list.append(value_node.id) elif isinstance(value_node, (ast.List, ast.Tuple)): for elt in value_node.elts: - if isinstance(elt, ast.Constant) and isinstance( - elt.value, str): + if isinstance(elt, ast.Constant) and isinstance(elt.value, str): model_cls_name_list.append(elt.value) elif isinstance(elt, ast.Name): model_cls_name_list.append(elt.id) @@ -229,14 +212,11 @@ def _discover_and_register_models() -> dict[str, tuple[str, str, str]]: mod_relname = ".".join(sub_parts) else: mod_base = filename[:-3] - mod_relname = ".".join(sub_parts + - [mod_base]) if sub_parts else mod_base + mod_relname = ".".join(sub_parts + [mod_base]) if sub_parts else mod_base for model_cls_str in model_cls_name_list: if model_cls_str in discovered_models: - logger.warning( - "Duplicate architecture found: %s. Overwriting.", - model_cls_str) + logger.warning("Duplicate architecture found: %s. Overwriting.", model_cls_str) discovered_models[model_cls_str] = ( component_name, mod_relname, @@ -244,8 +224,7 @@ def _discover_and_register_models() -> dict[str, tuple[str, str, str]]: ) except Exception as e: - logger.warning("Could not parse %s to find models: %s", - filepath, e) + logger.warning("Could not parse %s to find models: %s", filepath, e) return discovered_models @@ -316,9 +295,7 @@ def _run_in_subprocess(fn: Callable[[], _T]) -> _T: # cannot use `sys.executable __file__` here because the script # contains relative imports - returned = subprocess.run(_SUBPROCESS_COMMAND, - input=input_bytes, - capture_output=True) + returned = subprocess.run(_SUBPROCESS_COMMAND, input=input_bytes, capture_output=True) # check if the subprocess is successful try: @@ -343,8 +320,7 @@ class _LazyRegisteredModel(_BaseRegisteredModel): # Performed in another process to avoid initializing CUDA def inspect_model_cls(self) -> _ModelInfo: - return _run_in_subprocess( - lambda: _ModelInfo.from_model_cls(self.load_model_cls())) + return _run_in_subprocess(lambda: _ModelInfo.from_model_cls(self.load_model_cls())) def load_model_cls(self) -> type[nn.Module]: mod = importlib.import_module(self.module_name) @@ -373,8 +349,7 @@ def _try_inspect_model_cls( try: return model.inspect_model_cls() except Exception: - logger.exception("Error in inspecting model architecture '%s'", - model_arch) + logger.exception("Error in inspecting model architecture '%s'", model_arch) return None @@ -422,13 +397,11 @@ class _ModelRegistry: def _raise_for_unsupported(self, architectures: list[str]) -> NoReturn: all_supported_archs = self.get_supported_archs() if any(arch in all_supported_archs for arch in architectures): - raise ValueError( - f"Model architectures {architectures} failed " - "to be inspected. Please check the logs for more details.") + raise ValueError(f"Model architectures {architectures} failed " + "to be inspected. Please check the logs for more details.") - raise ValueError( - f"Model architectures {architectures} are not supported for now. " - f"Supported architectures: {all_supported_archs}") + raise ValueError(f"Model architectures {architectures} are not supported for now. " + f"Supported architectures: {all_supported_archs}") def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None: if model_arch not in self.models: @@ -493,6 +466,5 @@ ModelRegistry = _ModelRegistry({ component_name=component_name, class_name=cls_name, ) - for model_arch, (component_name, mod_relname, - cls_name) in _FAST_VIDEO_MODELS.items() + for model_arch, (component_name, mod_relname, cls_name) in _FAST_VIDEO_MODELS.items() }) diff --git a/fastvideo/models/schedulers/base.py b/fastvideo/models/schedulers/base.py index 4c647efc4..758cce334 100644 --- a/fastvideo/models/schedulers/base.py +++ b/fastvideo/models/schedulers/base.py @@ -16,9 +16,7 @@ class BaseScheduler(ABC): for attr in required_attributes: if not hasattr(self, attr): - raise AttributeError( - f"Subclasses of BaseScheduler must define '{attr}' property" - ) + raise AttributeError(f"Subclasses of BaseScheduler must define '{attr}' property") @abstractmethod def set_shift(self, shift: float) -> None: @@ -29,7 +27,5 @@ class BaseScheduler(ABC): pass @abstractmethod - def scale_model_input(self, - sample: torch.Tensor, - timestep: int | None = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: pass diff --git a/fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py index 95f6e808e..2108fc74c 100644 --- a/fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py +++ b/fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py @@ -22,7 +22,6 @@ import torch from fastvideo.models.schedulers.base import BaseScheduler - WeightType = Literal["uniform", "gaussian", "beta08"] @@ -112,7 +111,7 @@ class FlowMapEulerDiscreteScheduler(BaseScheduler): if weight_type == "uniform": w = torch.ones_like(t_norm) elif weight_type == "gaussian": - w = torch.exp(-0.5 * ((t_norm - 0.5) / 0.2) ** 2) + w = torch.exp(-0.5 * ((t_norm - 0.5) / 0.2)**2) elif weight_type == "beta08": w = t_norm.pow(1.0) * (1.0 - t_norm).clamp_min(0.0).pow(0.5) else: @@ -140,24 +139,19 @@ class FlowMapEulerDiscreteScheduler(BaseScheduler): AnyFlow paper's hand-tuned ``[999, 937, 833, 624, 0]`` schedule. """ if num_inference_steps <= 0: - raise ValueError( - "num_inference_steps must be positive, " - f"got {num_inference_steps}") + raise ValueError("num_inference_steps must be positive, " + f"got {num_inference_steps}") device = torch.device(device) if custom_timesteps is not None: - ts = torch.as_tensor( - custom_timesteps, dtype=torch.float32, device=device) + ts = torch.as_tensor(custom_timesteps, dtype=torch.float32, device=device) if ts.ndim != 1: - raise ValueError( - "custom_timesteps must be 1-D, got shape " - f"{tuple(ts.shape)}") + raise ValueError("custom_timesteps must be 1-D, got shape " + f"{tuple(ts.shape)}") if not torch.all(ts[:-1] >= ts[1:]): - raise ValueError( - "custom_timesteps must be descending (largest first)") + raise ValueError("custom_timesteps must be descending (largest first)") else: - ts_norm = torch.linspace( - 1.0, 0.0, num_inference_steps + 1, device=device) + ts_norm = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device) ts_norm = self.apply_shift(ts_norm) ts = ts_norm * self.num_train_timesteps @@ -194,9 +188,7 @@ class FlowMapEulerDiscreteScheduler(BaseScheduler): """Linear flow-matching interpolation: ``x_t = (1 - σ) * x_0 + σ * ε``, where ``σ = t / num_train_timesteps``. """ - sigma = (timestep.to(original_samples.device, - dtype=original_samples.dtype) - / float(self.num_train_timesteps)) + sigma = (timestep.to(original_samples.device, dtype=original_samples.dtype) / float(self.num_train_timesteps)) view: list[int] = [-1] + [1] * (original_samples.ndim - 1) sigma = sigma.view(*view) return (1.0 - sigma) * original_samples + sigma * noise diff --git a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py index b4f1935e1..5daffcdb1 100644 --- a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -49,8 +49,7 @@ class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput): prev_sample: torch.FloatTensor -class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, - BaseScheduler): +class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """ Euler scheduler. @@ -132,21 +131,14 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, sigma_data: float | None = None, use_reference_discrete_timesteps: bool = False, ): - if sum([ - self.config.use_beta_sigmas, self.config.use_exponential_sigmas, - self.config.use_karras_sigmas - ]) > 1: + if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1: raise ValueError( "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." ) if time_shift_type not in {"exponential", "linear"}: - raise ValueError( - "`time_shift_type` must either be 'exponential' or 'linear'.") + raise ValueError("`time_shift_type` must either be 'exponential' or 'linear'.") - timesteps = np.linspace(1, - num_train_timesteps, - num_train_timesteps, - dtype=np.float32)[::-1].copy() + timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) sigmas = timesteps / num_train_timesteps @@ -162,8 +154,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, self._shift = shift - self.sigmas = sigmas.to( - "cpu") # to avoid too much CPU/GPU communication + self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication self.sigma_min = sigma_min if sigma_min is not None else self.sigmas[-1].item() self.sigma_max = self.sigmas[0].item() @@ -228,8 +219,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, if sample.device.type == "mps" and torch.is_floating_point(timestep): # mps does not support float64 - schedule_timesteps = self.timesteps.to(sample.device, - dtype=torch.float32) + schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) assert isinstance(timestep, torch.Tensor) timestep = timestep.to(sample.device, dtype=torch.float32) else: @@ -239,9 +229,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index if self.begin_index is None: - step_indices = [ - self.index_for_timestep(t, schedule_timesteps) for t in timestep - ] + step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep] elif self.step_index is not None: # add_noise is called after first denoising step (for inpainting) step_indices = [self.step_index] * timestep.shape[0] @@ -260,15 +248,13 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, def _sigma_to_t(self, sigma: float) -> float: return sigma * self.config.num_train_timesteps - def time_shift(self, mu: float, sigma: float, - t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: + def time_shift(self, mu: float, sigma: float, t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: if self.config.time_shift_type == "exponential": return self._time_shift_exponential(mu, sigma, t) elif self.config.time_shift_type == "linear": return self._time_shift_linear(mu, sigma, t) else: - raise ValueError( - f"Unknown time_shift_type: {self.config.time_shift_type}") + raise ValueError(f"Unknown time_shift_type: {self.config.time_shift_type}") def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor: r""" @@ -318,19 +304,15 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, automatically. """ if self.config.use_dynamic_shifting and mu is None: - raise ValueError( - "`mu` must be passed when `use_dynamic_shifting` is set to be `True`" - ) + raise ValueError("`mu` must be passed when `use_dynamic_shifting` is set to be `True`") - if sigmas is not None and timesteps is not None and len(sigmas) != len( - timesteps): - raise ValueError( - "`sigmas` and `timesteps` should have the same length") + if sigmas is not None and timesteps is not None and len(sigmas) != len(timesteps): + raise ValueError("`sigmas` and `timesteps` should have the same length") if num_inference_steps is not None: - if (sigmas is not None and len(sigmas) != num_inference_steps) or ( - timesteps is not None - and len(timesteps) != num_inference_steps): + if (sigmas is not None + and len(sigmas) != num_inference_steps) or (timesteps is not None + and len(timesteps) != num_inference_steps): raise ValueError( "`sigmas` and `timesteps` should have the same length as num_inference_steps, if `num_inference_steps` is provided" ) @@ -340,9 +322,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, elif timesteps is not None: num_inference_steps = len(timesteps) else: - raise ValueError( - "Either num_inference_steps, sigmas, or timesteps must be provided" - ) + raise ValueError("Either num_inference_steps, sigmas, or timesteps must be provided") self.num_inference_steps = num_inference_steps @@ -383,48 +363,36 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, assert mu is not None, "mu cannot be None when use_dynamic_shifting is True" sigmas_array = self.time_shift(mu, 1.0, sigmas_array) else: - sigmas_array = self.shift * sigmas_array / ( - 1 + (self.shift - 1) * sigmas_array) + sigmas_array = self.shift * sigmas_array / (1 + (self.shift - 1) * sigmas_array) # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value if self.config.shift_terminal: - sigmas_tensor = torch.from_numpy(sigmas_array).to( - dtype=torch.float32) + sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32) sigmas_tensor = self.stretch_shift_to_terminal(sigmas_tensor) sigmas_array = sigmas_tensor.numpy() # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules if self.config.use_karras_sigmas: - sigmas_tensor = torch.from_numpy(sigmas_array).to( - dtype=torch.float32) - sigmas_tensor = self._convert_to_karras( - in_sigmas=sigmas_tensor, - num_inference_steps=num_inference_steps) + sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32) + sigmas_tensor = self._convert_to_karras(in_sigmas=sigmas_tensor, num_inference_steps=num_inference_steps) sigmas_array = sigmas_tensor.numpy() elif self.config.use_exponential_sigmas: - sigmas_tensor = torch.from_numpy(sigmas_array).to( - dtype=torch.float32) - sigmas_tensor = self._convert_to_exponential( - in_sigmas=sigmas_tensor, - num_inference_steps=num_inference_steps) + sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32) + sigmas_tensor = self._convert_to_exponential(in_sigmas=sigmas_tensor, + num_inference_steps=num_inference_steps) sigmas_array = sigmas_tensor.numpy() elif self.config.use_beta_sigmas: - sigmas_tensor = torch.from_numpy(sigmas_array).to( - dtype=torch.float32) - sigmas_tensor = self._convert_to_beta( - in_sigmas=sigmas_tensor, - num_inference_steps=num_inference_steps) + sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32) + sigmas_tensor = self._convert_to_beta(in_sigmas=sigmas_tensor, num_inference_steps=num_inference_steps) sigmas_array = sigmas_tensor.numpy() # 5. Convert sigmas and timesteps to tensors and move to specified device - sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32, - device=device) + sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32, device=device) if not is_timesteps_provided: timesteps_tensor = sigmas_tensor * self.config.num_train_timesteps else: assert timesteps_array is not None - timesteps_tensor = torch.from_numpy(timesteps_array).to( - dtype=torch.float32, device=device) + timesteps_tensor = torch.from_numpy(timesteps_array).to(dtype=torch.float32, device=device) # 6. Append the terminal sigma value. # If a model requires inverted sigma schedule for denoising but timesteps without inversion, the @@ -432,9 +400,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, if self.config.invert_sigmas: sigmas_tensor = 1.0 - sigmas_tensor timesteps_tensor = sigmas_tensor * self.config.num_train_timesteps - sigmas_tensor = torch.cat( - [sigmas_tensor, - torch.ones(1, device=sigmas_tensor.device)]) + sigmas_tensor = torch.cat([sigmas_tensor, torch.ones(1, device=sigmas_tensor.device)]) else: sigmas_tensor = torch.cat([sigmas_tensor, torch.zeros(1, device=sigmas_tensor.device)]) @@ -443,10 +409,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, self._step_index = None self._begin_index = None - def index_for_timestep( - self, - timestep: float | torch.FloatTensor, - schedule_timesteps: torch.Tensor | None = None) -> int: + def index_for_timestep(self, + timestep: float | torch.FloatTensor, + schedule_timesteps: torch.Tensor | None = None) -> int: if schedule_timesteps is None: schedule_timesteps = self.timesteps @@ -513,10 +478,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, """ if (isinstance(timestep, int | torch.IntTensor | torch.LongTensor)): - raise ValueError(( - "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" - " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass" - " one of the `scheduler.timesteps` as a timestep."), ) + raise ValueError(("Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep."), ) if self.step_index is None: self._init_step_index(timestep) @@ -566,8 +530,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample) # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras - def _convert_to_karras(self, in_sigmas: torch.Tensor, - num_inference_steps: int) -> torch.Tensor: + def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: """Constructs the noise schedule of Karras et al. (2022).""" # Hack to make sure that other schedulers which copy this function don't break @@ -587,8 +550,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, return torch.from_numpy(sigmas).to(dtype=in_sigmas.dtype, device=in_sigmas.device) # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, - num_inference_steps: int) -> torch.Tensor: + def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: """Constructs an exponential noise schedule.""" # Hack to make sure that other schedulers which copy this function don't break @@ -600,9 +562,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - sigmas = np.exp( - np.linspace(math.log(sigma_max), math.log(sigma_min), - num_inference_steps)) + sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) return torch.from_numpy(sigmas).to(dtype=in_sigmas.dtype, device=in_sigmas.device) # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta @@ -623,24 +583,19 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() sigmas = np.array([ - sigma_min + (ppf * (sigma_max - sigma_min)) for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] + sigma_min + (ppf * (sigma_max - sigma_min)) for ppf in + [scipy.stats.beta.ppf(timestep, alpha, beta) for timestep in 1 - np.linspace(0, 1, num_inference_steps)] ]) return torch.from_numpy(sigmas).to(dtype=in_sigmas.dtype, device=in_sigmas.device) - def _time_shift_exponential( - self, mu: float, sigma: float, - t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: + def _time_shift_exponential(self, mu: float, sigma: float, + t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: if isinstance(t, np.ndarray): return np.exp(mu) / (np.exp(mu) + (1 / t - 1)**sigma) else: return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma) - def _time_shift_linear( - self, mu: float, sigma: float, - t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: + def _time_shift_linear(self, mu: float, sigma: float, t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: return mu / (mu + (1 / t - 1)**sigma) def add_noise( @@ -649,7 +604,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, noise: torch.Tensor, timestep: torch.IntTensor, ) -> torch.Tensor: - """ Args: clean_latent: the clean latent with shape [B, C, H, W], @@ -675,19 +629,17 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, # timestep shape should be [B] self.sigmas = self.sigmas.to(noise.device) self.timesteps = self.timesteps.to(noise.device) - timestep_id = torch.argmin( - (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + timestep_id = torch.argmin((self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) sample = (1 - sigma) * clean_latent + sigma * noise return sample.type_as(noise) - def scale_model_input(self, - sample: torch.Tensor, - timestep: int | None = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: return sample def __len__(self) -> int: return self.config.num_train_timesteps + # Entry point for model registry EntryClass = FlowMatchEulerDiscreteScheduler diff --git a/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py b/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py index a5be46a76..b03cb88a2 100644 --- a/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py +++ b/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py @@ -9,9 +9,7 @@ from typing import Any import numpy as np import torch from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers, - SchedulerMixin, - SchedulerOutput) +from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput) from diffusers.utils import deprecate from fastvideo.models.schedulers.base import BaseScheduler @@ -98,14 +96,12 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): if solver_type in ["midpoint", "heun", "logrho"]: self.register_to_config(solver_type="bh2") else: - raise NotImplementedError( - f"{solver_type} is not implemented for {self.__class__}") + raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") self.predict_x0 = predict_x0 # setable values self.num_inference_steps: int | None = None - alphas = np.linspace(1, 1 / num_train_timesteps, - num_train_timesteps)[::-1].copy() + alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps)[::-1].copy() sigmas = 1.0 - alphas sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32) @@ -115,8 +111,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): if not use_dynamic_shifting: # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution assert shift is not None - sigmas = shift * sigmas / (1 + - (shift - 1) * sigmas) # pyright: ignore + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) # pyright: ignore self.sigmas = sigmas self.timesteps = sigmas * num_train_timesteps @@ -131,8 +126,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): self._step_index: int | None = None self._begin_index: int | None = None - self.sigmas = self.sigmas.to( - "cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication self.sigma_min = self.sigmas[-1].item() self.sigma_max = self.sigmas[0].item() @@ -187,10 +181,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """ if self.config.use_dynamic_shifting and mu is None: - raise ValueError( - " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`" - ) - + raise ValueError(" you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`") # Cosmos official uses `use_kerras_sigma=True` and a specific EDM sigma schedule. # Some external code uses the misspelling `use_kerras_sigma`; support both. @@ -206,8 +197,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): # Match the official Cosmos implementation: Karras/EDM schedule with # `num_inference_steps + 1` points, then `final_sigmas_type="zero"` # appends the terminal sigma (0.0). - ramp = np.arange(num_inference_steps + 1, - dtype=np.float32) / float(num_inference_steps) + ramp = np.arange(num_inference_steps + 1, dtype=np.float32) / float(num_inference_steps) min_inv_rho = sigma_min**(1 / rho) max_inv_rho = sigma_max**(1 / rho) sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho @@ -217,8 +207,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): if sigmas is None: assert num_inference_steps is not None sigmas = np.linspace(self.sigma_max, self.sigma_min, - num_inference_steps + - 1).copy()[:-1] # pyright: ignore + num_inference_steps + 1).copy()[:-1] # pyright: ignore if self.config.use_dynamic_shifting: assert mu is not None @@ -231,22 +220,18 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) # pyright: ignore if self.config.final_sigmas_type == "sigma_min": - sigma_last = ((1 - self.alphas_cumprod[0]) / - self.alphas_cumprod[0])**0.5 + sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0])**0.5 elif self.config.final_sigmas_type == "zero": sigma_last = 0 else: raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}") timesteps = sigmas * self.config.num_train_timesteps - sigmas = np.concatenate([sigmas, [sigma_last] - ]).astype(np.float32) # pyright: ignore + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) # pyright: ignore self.sigmas = torch.from_numpy(sigmas) - self.timesteps = torch.from_numpy(timesteps).to(device=device, - dtype=torch.int64) + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) self.num_inference_steps = len(timesteps) @@ -261,8 +246,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): # add an index counter for schedulers that allow duplicated timesteps self._step_index = None self._begin_index = None - self.sigmas = self.sigmas.to( - "cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: @@ -279,25 +263,19 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): batch_size, channels, *remaining_dims = sample.shape if dtype not in (torch.float32, torch.float64): - sample = sample.float( - ) # upcast for quantile calculation, and clamp not implemented for cpu half + sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half # Flatten sample for doing quantile calculation along each image sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) abs_sample = sample.abs() # "a certain percentile absolute pixel value" - s = torch.quantile(abs_sample, - self.config.dynamic_thresholding_ratio, - dim=1) + s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze( - 1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp( - sample, -s, s - ) / s # "we threshold xt0 to the range [-s, s] and then divide by s" + s, min=1, + max=self.config.sample_max_value) # When clamped to min=1, equivalent to standard clipping to [-1, 1] + s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 + sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" sample = sample.reshape(batch_size, channels, *remaining_dims) sample = sample.to(dtype) @@ -342,8 +320,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): if len(args) > 1: sample = args[1] else: - raise ValueError( - "missing `sample` as a required keyword argument") + raise ValueError("missing `sample` as a required keyword argument") if timestep is not None: deprecate( "timesteps", @@ -361,8 +338,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): else: raise ValueError( f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," - " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." - ) + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler.") if self.config.thresholding: x0_pred = self._threshold_sample(x0_pred) @@ -375,8 +351,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): else: raise ValueError( f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," - " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." - ) + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler.") if self.config.thresholding: sigma_t = self.sigmas[self.step_index] @@ -411,20 +386,17 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): `torch.Tensor`: The sample tensor at the previous timestep. """ - prev_timestep = args[0] if len(args) > 0 else kwargs.pop( - "prev_timestep", None) + prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None) if sample is None: if len(args) > 1: sample = args[1] else: - raise ValueError( - " missing `sample` as a required keyword argument") + raise ValueError(" missing `sample` as a required keyword argument") if order is None: if len(args) > 2: order = args[2] else: - raise ValueError( - " missing `order` as a required keyword argument") + raise ValueError(" missing `order` as a required keyword argument") if prev_timestep is not None: deprecate( "prev_timestep", @@ -441,17 +413,14 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): x_t = self.solver_p.step(model_output, s0, x).prev_sample return x_t - sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[ - self.step_index] # pyright: ignore + sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] # pyright: ignore alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) # Numerical safety eps = 1e-12 - lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log( - torch.clamp(sigma_t, min=eps)) - lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log( - torch.clamp(sigma_s0, min=eps)) + lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log(torch.clamp(sigma_t, min=eps)) + lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log(torch.clamp(sigma_s0, min=eps)) h = lambda_t - lambda_s0 device = sample.device @@ -462,8 +431,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): si = self.step_index - i # pyright: ignore mi = model_output_list[-(i + 1)] alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) - lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log( - torch.clamp(sigma_si, min=eps)) + lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log(torch.clamp(sigma_si, min=eps)) rk = (lambda_si - lambda_s0) / h rks.append(rk) assert mi is not None @@ -504,8 +472,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) else: assert isinstance(R, torch.Tensor) - rhos_p = torch.linalg.solve(R[:-1, :-1], - b[:-1]).to(device).to(x.dtype) + rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype) else: D1s = None @@ -549,26 +516,22 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): `torch.Tensor`: The corrected sample tensor at the current timestep. """ - this_timestep = args[0] if len(args) > 0 else kwargs.pop( - "this_timestep", None) + this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None) if last_sample is None: if len(args) > 1: last_sample = args[1] else: - raise ValueError( - " missing`last_sample` as a required keyword argument") + raise ValueError(" missing`last_sample` as a required keyword argument") if this_sample is None: if len(args) > 2: this_sample = args[2] else: - raise ValueError( - " missing`this_sample` as a required keyword argument") + raise ValueError(" missing`this_sample` as a required keyword argument") if order is None: if len(args) > 3: order = args[3] else: - raise ValueError( - " missing`order` as a required keyword argument") + raise ValueError(" missing`order` as a required keyword argument") if this_timestep is not None: deprecate( "this_timestep", @@ -583,16 +546,13 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): x_t = this_sample model_t = this_model_output - sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[ - self.step_index - 1] # pyright: ignore + sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1] # pyright: ignore alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) eps = 1e-12 - lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log( - torch.clamp(sigma_t, min=eps)) - lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log( - torch.clamp(sigma_s0, min=eps)) + lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log(torch.clamp(sigma_t, min=eps)) + lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log(torch.clamp(sigma_s0, min=eps)) h = lambda_t - lambda_s0 device = this_sample.device @@ -603,8 +563,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): si = self.step_index - (i + 1) # pyright: ignore mi = model_output_list[-(i + 1)] alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) - lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log( - torch.clamp(sigma_si, min=eps)) + lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log(torch.clamp(sigma_si, min=eps)) rk = (lambda_si - lambda_s0) / h rks.append(rk) assert mi is not None @@ -715,20 +674,17 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """ if self.num_inference_steps is None: raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler") if self.step_index is None: self._init_step_index(timestep) use_corrector = ( - self.step_index > 0 - and self.step_index - 1 not in self.disable_corrector + self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None # pyright: ignore ) - model_output_convert = self.convert_model_output(model_output, - sample=sample) + model_output_convert = self.convert_model_output(model_output, sample=sample) if use_corrector: sample = self.multistep_uni_c_bh_update( @@ -746,20 +702,16 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): self.timestep_list[-1] = timestep # pyright: ignore if self.config.lower_order_final: - this_order = min(self.config.solver_order, - len(self.timesteps) - - self.step_index) # pyright: ignore + this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index) # pyright: ignore else: this_order = self.config.solver_order - self.this_order: int = min(this_order, self.lower_order_nums + - 1) # warmup for multistep + self.this_order: int = min(this_order, self.lower_order_nums + 1) # warmup for multistep assert self.this_order > 0 self.last_sample = sample prev_sample = self.multistep_uni_p_bh_update( - model_output= - model_output, # pass the original non-converted model output, in case solver-p is used + model_output=model_output, # pass the original non-converted model output, in case solver-p is used sample=sample, order=self.this_order, ) @@ -776,8 +728,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): return SchedulerOutput(prev_sample=prev_sample) - def scale_model_input(self, sample: torch.Tensor, *args, - **kwargs) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -800,25 +751,18 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): timesteps: torch.IntTensor, ) -> torch.Tensor: # Make sure sigmas and timesteps have the same device and dtype as original_samples - sigmas = self.sigmas.to(device=original_samples.device, - dtype=original_samples.dtype) - if original_samples.device.type == "mps" and torch.is_floating_point( - timesteps): + sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) + if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, - dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, - dtype=torch.float32) + schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) + timesteps = timesteps.to(original_samples.device, dtype=torch.float32) else: schedule_timesteps = self.timesteps.to(original_samples.device) timesteps = timesteps.to(original_samples.device) # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index if self.begin_index is None: - step_indices = [ - self.index_for_timestep(t, schedule_timesteps) - for t in timesteps - ] + step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] elif self.step_index is not None: # add_noise is called after first denoising step (for inpainting) step_indices = [self.step_index] * timesteps.shape[0] @@ -837,5 +781,6 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): def __len__(self): return self.config.num_train_timesteps + # Entry point for model registry EntryClass = FlowUniPCMultistepScheduler diff --git a/fastvideo/models/schedulers/scheduling_minimax_h3.py b/fastvideo/models/schedulers/scheduling_minimax_h3.py index 44cf99ff7..647075caa 100644 --- a/fastvideo/models/schedulers/scheduling_minimax_h3.py +++ b/fastvideo/models/schedulers/scheduling_minimax_h3.py @@ -59,20 +59,15 @@ class MiniMaxH3Scheduler(SchedulerMixin, ConfigMixin): ) -> None: if sigmas is None: if num_inference_steps is None or num_inference_steps < 2: - raise ValueError( - "`set_timesteps` requires explicit `sigmas` or " - f"`num_inference_steps` >= 2, got {num_inference_steps}." - ) + raise ValueError("`set_timesteps` requires explicit `sigmas` or " + f"`num_inference_steps` >= 2, got {num_inference_steps}.") base = torch.linspace(1.0, 0.0, int(num_inference_steps), dtype=torch.float32) sigma_tensor = self._shift * base / (1 + (self._shift - 1) * base) sigma_tensor = torch.unique_consecutive(sigma_tensor) else: sigma_tensor = torch.as_tensor(sigmas, dtype=torch.float32).flatten().cpu() - is_valid = ( - sigma_tensor.numel() >= 2 - and bool((sigma_tensor[1:] < sigma_tensor[:-1]).all()) - and sigma_tensor[-1].item() == 0.0 - ) + is_valid = (sigma_tensor.numel() >= 2 and bool((sigma_tensor[1:] < sigma_tensor[:-1]).all()) + and sigma_tensor[-1].item() == 0.0) if not is_valid: raise ValueError("`sigmas` must hold at least two strictly decreasing values ending at 0.0.") @@ -140,7 +135,7 @@ class MiniMaxH3Scheduler(SchedulerMixin, ConfigMixin): self._step_index += 1 if not return_dict: - return (prev_sample,) + return (prev_sample, ) return MiniMaxH3SchedulerOutput(prev_sample=prev_sample) diff --git a/fastvideo/models/schedulers/scheduling_rcm.py b/fastvideo/models/schedulers/scheduling_rcm.py index 71cc714b0..e48066af7 100644 --- a/fastvideo/models/schedulers/scheduling_rcm.py +++ b/fastvideo/models/schedulers/scheduling_rcm.py @@ -71,17 +71,17 @@ class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): if mid_timesteps is None: mid_timesteps = [1.5, 1.4, 1.0] self._mid_timesteps = mid_timesteps - + self.num_train_timesteps = num_train_timesteps self.sigma_max = sigma_max - + # Initialize with default timesteps (will be set properly via set_timesteps) self.timesteps = torch.tensor([1.0, 0.0], dtype=torch.float64) self.sigmas = self.timesteps.clone() - + self._step_index: int | None = None self._begin_index: int | None = None - + BaseScheduler.__init__(self) @property @@ -148,13 +148,11 @@ class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): Override the initial sigma value. """ if num_inference_steps < 1 or num_inference_steps > 4: - logger.warning( - "rCM is optimized for 1-4 steps, got %d steps. " - "Performance may be suboptimal.", num_inference_steps - ) + logger.warning("rCM is optimized for 1-4 steps, got %d steps. " + "Performance may be suboptimal.", num_inference_steps) self.num_inference_steps = num_inference_steps - + if sigma_max is not None: self.sigma_max = sigma_max @@ -173,10 +171,10 @@ class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): # Store raw sigmas for use in step() formula self.sigmas = t_steps.clone() - + # Scale timesteps by 1000 for model input (as per TurboDiffusion) self.timesteps = t_steps * 1000 - + self._step_index = None self._begin_index = None @@ -222,7 +220,7 @@ class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """ if noise is None: raise ValueError("noise must be provided for rCM scale_noise") - + # Use raw sigma (not scaled timestep) for initial noise scaling t_initial = self.sigmas[0] return noise.to(torch.float64) * t_initial @@ -262,7 +260,7 @@ class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): self._init_step_index() assert self._step_index is not None - + # Get current and next sigma values (raw, unscaled) for rCM formula # Note: self.timesteps is scaled by 1000 for model input, # but we need raw values for the step formula @@ -279,7 +277,7 @@ class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): # rCM update: x = (1 - t_next) * (x - t_cur * v_pred) + t_next * noise x_denoised = sample - t_cur * model_output - + # Generate noise for SDE sampling if isinstance(generator, list): generator = generator[0] @@ -299,7 +297,7 @@ class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): prev_sample = prev_sample.to(model_output.dtype) if not return_dict: - return (prev_sample,) + return (prev_sample, ) return RCMSchedulerOutput(prev_sample=prev_sample) @@ -314,13 +312,12 @@ class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): Not typically used for rCM inference, but provided for API compatibility. """ - raise NotImplementedError( - "add_noise is not implemented for RCMScheduler. " - "Use scale_noise for initializing the sampling process." - ) + raise NotImplementedError("add_noise is not implemented for RCMScheduler. " + "Use scale_noise for initializing the sampling process.") def __len__(self) -> int: return self.config.num_train_timesteps + # Entry point for model registry EntryClass = RCMScheduler diff --git a/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py b/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py index b96118b31..915bd8809 100644 --- a/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py +++ b/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py @@ -21,12 +21,23 @@ class SelfForcingFlowMatchSchedulerOutput(BaseOutput): """ prev_sample: torch.FloatTensor + class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): - + config_name = "scheduler_config.json" order = 1 + @register_to_config - def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False, training=False): + def __init__(self, + num_inference_steps=100, + num_train_timesteps=1000, + shift=3.0, + sigma_max=1.0, + sigma_min=0.003 / 1.002, + inverse_timesteps=False, + extra_one_step=False, + reverse_sigmas=False, + training=False): self.num_train_timesteps = num_train_timesteps self.shift = shift self.sigma_max = sigma_max @@ -36,15 +47,18 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): self.reverse_sigmas = reverse_sigmas self.set_timesteps(num_inference_steps, training=training) - def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False, return_dict=False, **kwargs): + def set_timesteps(self, + num_inference_steps=100, + denoising_strength=1.0, + training=False, + return_dict=False, + **kwargs): sigma_start = self.sigma_min + \ (self.sigma_max - self.sigma_min) * denoising_strength if self.extra_one_step: - self.sigmas = torch.linspace( - sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] else: - self.sigmas = torch.linspace( - sigma_start, self.sigma_min, num_inference_steps) + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps) if self.inverse_timesteps: self.sigmas = torch.flip(self.sigmas, dims=[0]) self.sigmas = self.shift * self.sigmas / \ @@ -54,14 +68,19 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): self.timesteps = self.sigmas * self.num_train_timesteps if training: x = self.timesteps - y = torch.exp(-2 * ((x - num_inference_steps / 2) / - num_inference_steps) ** 2) + y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps)**2) y_shifted = y - y.min() bsmntw_weighing = y_shifted * \ (num_inference_steps / y_shifted.sum()) self.linear_timesteps_weights = bsmntw_weighing - def step(self, model_output: torch.FloatTensor, timestep: torch.FloatTensor, sample: torch.FloatTensor, to_final=False, return_dict=False, **kwargs): + def step(self, + model_output: torch.FloatTensor, + timestep: torch.FloatTensor, + sample: torch.FloatTensor, + to_final=False, + return_dict=False, + **kwargs): if timestep.ndim == 2: timestep = timestep.flatten(0, 1) elif timestep.ndim == 0: @@ -73,12 +92,10 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): self.timesteps = self.timesteps.to(model_output.device) timestep = timestep.to(model_output.device) - timestep_id = torch.argmin( - (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + timestep_id = torch.argmin((self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) if to_final or (timestep_id + 1 >= len(self.timesteps)).any(): - sigma_ = 1 if ( - self.inverse_timesteps or self.reverse_sigmas) else 0 + sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 else: sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1) prev_sample = sample + model_output * (sigma_ - sigma) @@ -89,9 +106,9 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): @staticmethod def calculate_alpha_beta_high(sigma, sigma_bound): alpha = (1 - sigma) / (1 - sigma_bound) - beta = torch.sqrt(sigma ** 2 - (alpha * sigma_bound) ** 2) + beta = torch.sqrt(sigma**2 - (alpha * sigma_bound)**2) return alpha, beta - + def add_noise(self, original_samples, noise, timestep): """ Diffusion forward corruption process. @@ -105,8 +122,7 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): timestep = timestep.flatten(0, 1) self.sigmas = self.sigmas.to(noise.device) self.timesteps = self.timesteps.to(noise.device) - timestep_id = torch.argmin( - (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + timestep_id = torch.argmin((self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) sample = (1 - sigma) * original_samples + sigma * noise return sample.type_as(noise) @@ -126,12 +142,10 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): boundary_timestep = boundary_timestep.flatten(0, 1) self.sigmas = self.sigmas.to(noise.device) self.timesteps = self.timesteps.to(noise.device) - timestep_id = torch.argmin( - (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + timestep_id = torch.argmin((self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) - boundary_timestep_id = torch.argmin( - (self.timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1) + boundary_timestep_id = torch.argmin((self.timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1) sigma_boundary = self.sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1) alpha, beta = self.calculate_alpha_beta_high(sigma, sigma_boundary) sample = alpha * original_samples + beta * noise @@ -150,16 +164,16 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): if timestep.ndim == 2: timestep = timestep.flatten(0, 1) self.linear_timesteps_weights = self.linear_timesteps_weights.to(timestep.device) - timestep_id = torch.argmin( - (self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0) + timestep_id = torch.argmin((self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0) weights = self.linear_timesteps_weights[timestep_id] return weights - + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: return sample def set_shift(self, shift: float) -> None: self.shift = shift + # Entry point for model registry EntryClass = SelfForcingFlowMatchScheduler diff --git a/fastvideo/models/schedulers/scheduling_unipc_multistep.py b/fastvideo/models/schedulers/scheduling_unipc_multistep.py index 7d0c5b2ce..0b1e47f88 100644 --- a/fastvideo/models/schedulers/scheduling_unipc_multistep.py +++ b/fastvideo/models/schedulers/scheduling_unipc_multistep.py @@ -63,7 +63,7 @@ def betas_for_alpha_bar( if alpha_transform_type == "cosine": def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + return math.cos((t + 0.008) / 1.008 * math.pi / 2)**2 elif alpha_transform_type == "exp": @@ -235,7 +235,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) elif beta_schedule == "scaled_linear": # this schedule is very specific to the latent diffusion model. - self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2 + self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32)**2 elif beta_schedule == "squaredcos_cap_v2": # Glide cosine schedule self.betas = betas_for_alpha_bar(num_train_timesteps) @@ -257,7 +257,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): self.alpha_t = torch.sqrt(self.alphas_cumprod) self.sigma_t = torch.sqrt(1 - self.alphas_cumprod) self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t) - self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5 + self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod)**0.5 # standard deviation of the initial noise distribution self.init_noise_sigma = 1.0 @@ -314,9 +314,10 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """ self._begin_index = begin_index - def set_timesteps( - self, num_inference_steps: int, device: Union[str, torch.device] = None, mu: Optional[float] = None - ): + def set_timesteps(self, + num_inference_steps: int, + device: Union[str, torch.device] = None, + mu: Optional[float] = None): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -331,12 +332,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): assert self.config.use_dynamic_shifting and self.config.time_shift_type == "exponential" self.config.flow_shift = np.exp(mu) if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps + 1) - .round()[::-1][:-1] - .copy() - .astype(np.int64) - ) + timesteps = (np.linspace(0, self.config.num_train_timesteps - 1, + num_inference_steps + 1).round()[::-1][:-1].copy().astype(np.int64)) elif self.config.timestep_spacing == "leading": step_ratio = self.config.num_train_timesteps // (num_inference_steps + 1) # creates integer timesteps by multiplying by ratio @@ -354,7 +351,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'." ) - sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) + sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod)**0.5) if self.config.use_karras_sigmas: log_sigmas = np.log(sigmas) sigmas = np.flip(sigmas).copy() @@ -414,7 +411,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): else: sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) if self.config.final_sigmas_type == "sigma_min": - sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5 + sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0])**0.5 elif self.config.final_sigmas_type == "zero": sigma_last = 0 else: @@ -465,8 +462,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] + s, min=1, + max=self.config.sample_max_value) # When clamped to min=1, equivalent to standard clipping to [-1, 1] s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" @@ -505,7 +502,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): alpha_t = 1 - sigma sigma_t = sigma else: - alpha_t = 1 / ((sigma**2 + 1) ** 0.5) + alpha_t = 1 / ((sigma**2 + 1)**0.5) sigma_t = sigma * alpha_t return alpha_t, sigma_t @@ -525,9 +522,9 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): rho = 7.0 # 7.0 is the value used in the paper ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho + min_inv_rho = sigma_min**(1 / rho) + max_inv_rho = sigma_max**(1 / rho) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho return sigmas # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential @@ -547,9 +544,11 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): return sigmas # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: + def _convert_to_beta(self, + in_sigmas: torch.Tensor, + num_inference_steps: int, + alpha: float = 0.6, + beta: float = 0.6) -> torch.Tensor: """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" # Hack to make sure that other schedulers which copy this function don't break @@ -561,15 +560,10 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) + sigmas = np.array([ + sigma_min + (ppf * (sigma_max - sigma_min)) for ppf in + [scipy.stats.beta.ppf(timestep, alpha, beta) for timestep in 1 - np.linspace(0, 1, num_inference_steps)] + ]) return sigmas def convert_model_output( @@ -623,8 +617,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): else: raise ValueError( f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, " - "`v_prediction`, or `flow_prediction` for the UniPCMultistepScheduler." - ) + "`v_prediction`, or `flow_prediction` for the UniPCMultistepScheduler.") if self.config.thresholding: x0_pred = self._threshold_sample(x0_pred) @@ -642,8 +635,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): else: raise ValueError( f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction` for the UniPCMultistepScheduler." - ) + " `v_prediction` for the UniPCMultistepScheduler.") def multistep_uni_p_bh_update( self, @@ -958,15 +950,13 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """ if self.num_inference_steps is None: raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler") if self.step_index is None: self._init_step_index(timestep) - use_corrector = ( - self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None - ) + use_corrector = (self.step_index > 0 and self.step_index - 1 not in self.disable_corrector + and self.last_sample is not None) model_output_convert = self.convert_model_output(model_output, sample=sample) if use_corrector: @@ -1006,7 +996,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): self._step_index += 1 if not return_dict: - return (prev_sample,) + return (prev_sample, ) return SchedulerOutput(prev_sample=prev_sample) @@ -1063,5 +1053,6 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): def __len__(self): return self.config.num_train_timesteps + # Entry point for model registry EntryClass = UniPCMultistepScheduler diff --git a/fastvideo/models/upsamplers/hunyuan15.py b/fastvideo/models/upsamplers/hunyuan15.py index 0d31a517e..006cb0a41 100644 --- a/fastvideo/models/upsamplers/hunyuan15.py +++ b/fastvideo/models/upsamplers/hunyuan15.py @@ -32,7 +32,9 @@ from fastvideo.models.vaes.hunyuan15vae import ( from fastvideo.layers.activation import get_act_fn from fastvideo.configs.models.upsamplers import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig + class HunyuanVideo15ResnetBlock(nn.Module): + def __init__( self, in_channels: int, @@ -70,7 +72,9 @@ class HunyuanVideo15ResnetBlock(nn.Module): return hidden_states + residual + class SRResidualCausalBlock3D(nn.Module): + def __init__(self, channels: int): super().__init__() self.block = nn.Sequential( @@ -167,4 +171,4 @@ class SRTo1080pUpsampler(nn.Module): h = self.norm_out(h) h = get_act_fn("swish")(h) h = self.conv_out(h) - return h \ No newline at end of file + return h diff --git a/fastvideo/models/upsamplers/ltx2_upsampler.py b/fastvideo/models/upsamplers/ltx2_upsampler.py index c6bef8aba..8bf9f9902 100644 --- a/fastvideo/models/upsamplers/ltx2_upsampler.py +++ b/fastvideo/models/upsamplers/ltx2_upsampler.py @@ -211,8 +211,7 @@ class LatentUpsampler(nn.Module): raise ValueError("Either spatial_upsample or temporal_upsample must be True") self.post_upsample_res_blocks = nn.ModuleList( - [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] - ) + [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]) self.final_conv = conv(mid_channels, in_channels, kernel_size=3, padding=1) diff --git a/fastvideo/models/utils.py b/fastvideo/models/utils.py index 493baeb81..ba5898ba1 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/models/utils.py @@ -67,8 +67,7 @@ def set_weight_attrs( if weight_attrs is None: return for key, value in weight_attrs.items(): - assert not hasattr( - weight, key), (f"Overwriting existing tensor attribute: {key}") + assert not hasattr(weight, key), (f"Overwriting existing tensor attribute: {key}") # NOTE(woosuk): During weight loading, we often do something like: # narrowed_tensor = param.data.narrow(0, offset, len) @@ -115,9 +114,7 @@ def extract_layer_index(layer_name: str) -> int: return int_vals[0] -def modulate(x: torch.Tensor, - shift: torch.Tensor | None = None, - scale: torch.Tensor | None = None) -> torch.Tensor: +def modulate(x: torch.Tensor, shift: torch.Tensor | None = None, scale: torch.Tensor | None = None) -> torch.Tensor: """modulate by shift and scale Args: @@ -135,13 +132,10 @@ def modulate(x: torch.Tensor, elif scale is None: return x + shift.unsqueeze(1) # type: ignore[union-attr] else: - return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze( - 1) # type: ignore[union-attr] + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) # type: ignore[union-attr] -def pred_noise_to_pred_video(pred_noise: torch.Tensor, - noise_input_latent: torch.Tensor, - timestep: torch.Tensor, +def pred_noise_to_pred_video(pred_noise: torch.Tensor, noise_input_latent: torch.Tensor, timestep: torch.Tensor, scheduler: Any) -> torch.Tensor: """ Convert predicted noise to clean latent. @@ -175,17 +169,14 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor, noise_input_latent = noise_input_latent.double().to(device) sigmas = scheduler.sigmas.double().to(device) timesteps = scheduler.timesteps.double().to(device) - timestep_id = torch.argmin( - (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) pred_video = noise_input_latent - sigma_t * pred_noise return pred_video.to(dtype) -def pred_noise_to_x_bound(pred_noise: torch.Tensor, - noise_input_latent: torch.Tensor, - timestep: torch.Tensor, - boundary_timestep: torch.Tensor, - scheduler: Any) -> torch.Tensor: + +def pred_noise_to_x_bound(pred_noise: torch.Tensor, noise_input_latent: torch.Tensor, timestep: torch.Tensor, + boundary_timestep: torch.Tensor, scheduler: Any) -> torch.Tensor: """ Convert predicted noise to clean latent. @@ -219,12 +210,10 @@ def pred_noise_to_x_bound(pred_noise: torch.Tensor, noise_input_latent = noise_input_latent.double().to(device) sigmas = scheduler.sigmas.double().to(device) timesteps = scheduler.timesteps.double().to(device) - timestep_id = torch.argmin( - (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) - boundary_timestep_id = torch.argmin( - (timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1) + boundary_timestep_id = torch.argmin((timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1) sigma_t_boundary = sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1) pred_video = noise_input_latent - (sigma_t - sigma_t_boundary) * pred_noise - return pred_video.to(dtype) \ No newline at end of file + return pred_video.to(dtype) diff --git a/fastvideo/models/vaes/autoencoder_kl.py b/fastvideo/models/vaes/autoencoder_kl.py index d28f3e642..fd7c4c3a5 100644 --- a/fastvideo/models/vaes/autoencoder_kl.py +++ b/fastvideo/models/vaes/autoencoder_kl.py @@ -10,7 +10,7 @@ from fastvideo.configs.models import VAEConfig class AutoencoderKL(_DiffusersAutoencoderKL): - + def __init__(self, config: VAEConfig, **kwargs: Any) -> None: self.fastvideo_config = config arch = config.arch_config diff --git a/fastvideo/models/vaes/common.py b/fastvideo/models/vaes/common.py index 9ed7f0f5b..6cb85c982 100644 --- a/fastvideo/models/vaes/common.py +++ b/fastvideo/models/vaes/common.py @@ -64,13 +64,11 @@ class ParallelTiledVAE(ABC): def encode(self, x: torch.Tensor) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = x.shape - latent_num_frames = (num_frames - - 1) // self.temporal_compression_ratio + 1 + latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1 if self.use_tiling and self.use_temporal_tiling and num_frames > self.tile_sample_min_num_frames: latents = self.tiled_encode(x)[:, :, :latent_num_frames] - elif self.use_tiling and (width > self.tile_sample_min_width - or height > self.tile_sample_min_height): + elif self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): latents = self.spatial_tiled_encode(x)[:, :, :latent_num_frames] else: latents = self._encode(x)[:, :, :latent_num_frames] @@ -81,43 +79,38 @@ class ParallelTiledVAE(ABC): tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio - num_sample_frames = (num_frames - - 1) * self.temporal_compression_ratio + 1 + num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1 - if self.use_tiling and self.use_parallel_tiling and get_sp_world_size( - ) > 1: + if self.use_tiling and self.use_parallel_tiling and get_sp_world_size() > 1: return self.parallel_tiled_decode(z)[:, :, :num_sample_frames] if self.use_tiling and self.use_temporal_tiling and num_frames > tile_latent_min_num_frames: return self.tiled_decode(z)[:, :, :num_sample_frames] - if self.use_tiling and (width > tile_latent_min_width - or height > tile_latent_min_height): + if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height): return self.spatial_tiled_decode(z)[:, :, :num_sample_frames] return self._decode(z)[:, :, :num_sample_frames] - def blend_v(self, a: torch.Tensor, b: torch.Tensor, - blend_extent: int) -> torch.Tensor: + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) for y in range(blend_extent): - b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * ( - 1 - y / blend_extent) + b[:, :, :, y, :] * (y / blend_extent) + b[:, :, :, + y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (y / blend_extent) return b - def blend_h(self, a: torch.Tensor, b: torch.Tensor, - blend_extent: int) -> torch.Tensor: + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) for x in range(blend_extent): - b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * ( - 1 - x / blend_extent) + b[:, :, :, :, x] * (x / blend_extent) + b[:, :, :, :, + x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (x / blend_extent) return b - def blend_t(self, a: torch.Tensor, b: torch.Tensor, - blend_extent: int) -> torch.Tensor: + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) for x in range(blend_extent): - b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * ( - 1 - x / blend_extent) + b[:, :, x, :, :] * (x / blend_extent) + b[:, :, + x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, + x, :, :] * (x / blend_extent) return b def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor: @@ -148,26 +141,21 @@ class ParallelTiledVAE(ABC): for i in range(0, height, self.tile_sample_stride_height): row = [] for j in range(0, width, self.tile_sample_stride_width): - tile = x[:, :, :, i:i + self.tile_sample_min_height, - j:j + self.tile_sample_min_width] + tile = x[:, :, :, i:i + self.tile_sample_min_height, j:j + self.tile_sample_min_width] tile = self._encode(tile) row.append(tile) rows.append(row) - return self._merge_spatial_tiles(rows, blend_height, blend_width, - tile_latent_stride_height, + return self._merge_spatial_tiles(rows, blend_height, blend_width, tile_latent_stride_height, tile_latent_stride_width) - def _parallel_data_generator( - self, gathered_results, - gathered_dim_metadata) -> Iterator[tuple[torch.Tensor, int]]: + def _parallel_data_generator(self, gathered_results, gathered_dim_metadata) -> Iterator[tuple[torch.Tensor, int]]: global_idx = 0 for i, per_rank_metadata in enumerate(gathered_dim_metadata): _start_shape = 0 for shape in per_rank_metadata: mul_shape = prod(shape) - yield (gathered_results[i, _start_shape:_start_shape + - mul_shape].reshape(shape), global_idx) + yield (gathered_results[i, _start_shape:_start_shape + mul_shape].reshape(shape), global_idx) _start_shape += mul_shape global_idx += 1 @@ -190,12 +178,9 @@ class ParallelTiledVAE(ABC): blend_width = self.tile_sample_min_width - self.tile_sample_stride_width # Calculate tile dimensions - num_t_tiles = (T + tile_latent_stride_num_frames - - 1) // tile_latent_stride_num_frames - num_h_tiles = (H + tile_latent_stride_height - - 1) // tile_latent_stride_height - num_w_tiles = (W + tile_latent_stride_width - - 1) // tile_latent_stride_width + num_t_tiles = (T + tile_latent_stride_num_frames - 1) // tile_latent_stride_num_frames + num_h_tiles = (H + tile_latent_stride_height - 1) // tile_latent_stride_height + num_w_tiles = (W + tile_latent_stride_width - 1) // tile_latent_stride_width total_spatial_tiles = num_h_tiles * num_w_tiles total_tiles = num_t_tiles * total_spatial_tiles @@ -207,8 +192,7 @@ class ParallelTiledVAE(ABC): local_results = [] local_dim_metadata = [] # Process assigned tiles - for local_idx, global_idx in enumerate( - range(start_tile_idx, end_tile_idx)): + for local_idx, global_idx in enumerate(range(start_tile_idx, end_tile_idx)): t_idx = global_idx // total_spatial_tiles spatial_idx = global_idx % total_spatial_tiles h_idx = spatial_idx // num_w_tiles @@ -220,8 +204,7 @@ class ParallelTiledVAE(ABC): w_start = w_idx * tile_latent_stride_width # Extract and process tile - tile = z[:, :, t_start:t_start + tile_latent_min_num_frames + 1, - h_start:h_start + tile_latent_min_height, + tile = z[:, :, t_start:t_start + tile_latent_min_num_frames + 1, h_start:h_start + tile_latent_min_height, w_start:w_start + tile_latent_min_width] # Process tile @@ -240,13 +223,8 @@ class ParallelTiledVAE(ABC): results = torch.cat(local_results, dim=0).contiguous() del local_results # first gather size to pad the results - local_size = torch.tensor([results.size(0)], - device=results.device, - dtype=torch.int64) - all_sizes = [ - torch.zeros(1, device=results.device, dtype=torch.int64) - for _ in range(world_size) - ] + local_size = torch.tensor([results.size(0)], device=results.device, dtype=torch.int64) + all_sizes = [torch.zeros(1, device=results.device, dtype=torch.int64) for _ in range(world_size)] dist.all_gather(all_sizes, local_size) max_size = max(size.item() for size in all_sizes) padded_results = torch.zeros(max_size, device=results.device) @@ -256,17 +234,14 @@ class ParallelTiledVAE(ABC): # Gather all results gathered_dim_metadata = [None] * world_size gathered_results = torch.zeros_like(padded_results).repeat( - world_size, *[1] * len(padded_results.shape) - ).contiguous( - ) # use contiguous to make sure it won't copy data in the following operations + world_size, *[1] * len(padded_results.shape)).contiguous( + ) # use contiguous to make sure it won't copy data in the following operations # TODO (PY): use fastvideo distributed methods dist.all_gather_into_tensor(gathered_results, padded_results) dist.all_gather_object(gathered_dim_metadata, local_dim_metadata) # Process gathered results - data: list = [[[[] for _ in range(num_w_tiles)] - for _ in range(num_h_tiles)] for _ in range(num_t_tiles)] - for current_data, global_idx in self._parallel_data_generator( - gathered_results, gathered_dim_metadata): + data: list = [[[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)] for _ in range(num_t_tiles)] + for current_data, global_idx in self._parallel_data_generator(gathered_results, gathered_dim_metadata): t_idx = global_idx // total_spatial_tiles spatial_idx = global_idx % total_spatial_tiles h_idx = spatial_idx // num_w_tiles @@ -276,25 +251,19 @@ class ParallelTiledVAE(ABC): result_slices = [] last_slice_data = None for i, tem_data in enumerate(data): - slice_data = self._merge_spatial_tiles( - tem_data, blend_height, blend_width, - self.tile_sample_stride_height, self.tile_sample_stride_width) + slice_data = self._merge_spatial_tiles(tem_data, blend_height, blend_width, self.tile_sample_stride_height, + self.tile_sample_stride_width) if i > 0: - slice_data = self.blend_t(last_slice_data, slice_data, - self.blend_num_frames) - result_slices.append( - slice_data[:, :, :self.tile_sample_stride_num_frames, :, :]) + slice_data = self.blend_t(last_slice_data, slice_data, self.blend_num_frames) + result_slices.append(slice_data[:, :, :self.tile_sample_stride_num_frames, :, :]) else: - result_slices.append( - slice_data[:, :, :self.tile_sample_stride_num_frames + - 1, :, :]) + result_slices.append(slice_data[:, :, :self.tile_sample_stride_num_frames + 1, :, :]) last_slice_data = slice_data dec = torch.cat(result_slices, dim=2) return dec - def _merge_spatial_tiles(self, tiles, blend_height, blend_width, - stride_height, stride_width) -> torch.Tensor: + def _merge_spatial_tiles(self, tiles, blend_height, blend_width, stride_height, stride_width) -> torch.Tensor: """Helper function to merge spatial tiles with blending""" result_rows = [] for i, row in enumerate(tiles): @@ -338,13 +307,11 @@ class ParallelTiledVAE(ABC): for i in range(0, height, tile_latent_stride_height): row = [] for j in range(0, width, tile_latent_stride_width): - tile = z[:, :, :, i:i + tile_latent_min_height, - j:j + tile_latent_min_width] + tile = z[:, :, :, i:i + tile_latent_min_height, j:j + tile_latent_min_width] decoded = self._decode(tile) row.append(decoded) rows.append(row) - return self._merge_spatial_tiles(rows, blend_height, blend_width, - self.tile_sample_stride_height, + return self._merge_spatial_tiles(rows, blend_height, blend_width, self.tile_sample_stride_height, self.tile_sample_stride_width) def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: @@ -356,8 +323,7 @@ class ParallelTiledVAE(ABC): row = [] for i in range(0, num_frames, self.tile_sample_stride_num_frames): tile = x[:, :, i:i + self.tile_sample_min_num_frames + 1, :, :] - if self.use_tiling and (height > self.tile_sample_min_height - or width > self.tile_sample_min_width): + if self.use_tiling and (height > self.tile_sample_min_height or width > self.tile_sample_min_width): tile = self.spatial_tiled_encode(tile) else: tile = self._encode(tile) @@ -371,13 +337,10 @@ class ParallelTiledVAE(ABC): result_row = [] for i, tile in enumerate(row): if i > 0: - tile = self.blend_t(row[i - 1], tile, - blend_latent_num_frames) - result_row.append( - tile[:, :, :tile_latent_stride_num_frames, :, :]) + tile = self.blend_t(row[i - 1], tile, blend_latent_num_frames) + result_row.append(tile[:, :, :tile_latent_stride_num_frames, :, :]) else: - result_row.append(tile[:, :, :tile_latent_stride_num_frames + - 1, :, :]) + result_row.append(tile[:, :, :tile_latent_stride_num_frames + 1, :, :]) enc = torch.cat(result_row, dim=2) return enc @@ -392,8 +355,7 @@ class ParallelTiledVAE(ABC): row = [] for i in range(0, num_frames, tile_latent_stride_num_frames): tile = z[:, :, i:i + tile_latent_min_num_frames + 1, :, :] - if self.use_tiling and (tile.shape[-1] > tile_latent_min_width - or tile.shape[-2] > tile_latent_min_height): + if self.use_tiling and (tile.shape[-1] > tile_latent_min_width or tile.shape[-2] > tile_latent_min_height): decoded = self.spatial_tiled_decode(tile) else: decoded = self._decode(tile) @@ -404,11 +366,9 @@ class ParallelTiledVAE(ABC): for i, tile in enumerate(row): if i > 0: tile = self.blend_t(row[i - 1], tile, self.blend_num_frames) - result_row.append( - tile[:, :, :self.tile_sample_stride_num_frames, :, :]) + result_row.append(tile[:, :, :self.tile_sample_stride_num_frames, :, :]) else: - result_row.append( - tile[:, :, :self.tile_sample_stride_num_frames + 1, :, :]) + result_row.append(tile[:, :, :self.tile_sample_stride_num_frames + 1, :, :]) dec = torch.cat(result_row, dim=2) return dec @@ -483,10 +443,9 @@ class DiagonalGaussianDistribution: self.std = torch.exp(0.5 * self.logvar) self.var = torch.exp(self.logvar) if self.deterministic: - self.var = self.std = torch.zeros_like( - self.mean, - device=self.parameters.device, - dtype=self.parameters.dtype) + self.var = self.std = torch.zeros_like(self.mean, + device=self.parameters.device, + dtype=self.parameters.dtype) def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: # make sure sample is on the same device as the parameters and has same dtype @@ -499,9 +458,7 @@ class DiagonalGaussianDistribution: x = self.mean + self.std * sample return x - def kl(self, - other: Optional["DiagonalGaussianDistribution"] = None - ) -> torch.Tensor: + def kl(self, other: Optional["DiagonalGaussianDistribution"] = None) -> torch.Tensor: if self.deterministic: return torch.Tensor([0.0]) else: @@ -512,20 +469,17 @@ class DiagonalGaussianDistribution: ) else: return 0.5 * torch.sum( - torch.pow(self.mean - other.mean, 2) / other.var + - self.var / other.var - 1.0 - self.logvar + other.logvar, + torch.pow(self.mean - other.mean, 2) / other.var + self.var / other.var - 1.0 - self.logvar + + other.logvar, dim=[1, 2, 3], ) - def nll( - self, sample: torch.Tensor, - dims: tuple[int, ...] = (1, 2, 3)) -> torch.Tensor: + def nll(self, sample: torch.Tensor, dims: tuple[int, ...] = (1, 2, 3)) -> torch.Tensor: if self.deterministic: return torch.Tensor([0.0]) logtwopi = np.log(2.0 * np.pi) return 0.5 * torch.sum( - logtwopi + self.logvar + - torch.pow(sample - self.mean, 2) / self.var, + logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, dim=dims, ) diff --git a/fastvideo/models/vaes/cosmos25wanvae.py b/fastvideo/models/vaes/cosmos25wanvae.py index 4b1f1a0b0..3a113e03e 100644 --- a/fastvideo/models/vaes/cosmos25wanvae.py +++ b/fastvideo/models/vaes/cosmos25wanvae.py @@ -120,9 +120,7 @@ class Cosmos25WanVAEAdapter(nn.Module): z_mean = enc_out else: attrs = [a for a in dir(enc_out) if not a.startswith("_")] - raise RuntimeError( - f"Unsupported VAE encoder output type: {type(enc_out)}. attrs={attrs}" - ) + raise RuntimeError(f"Unsupported VAE encoder output type: {type(enc_out)}. attrs={attrs}") mean, std = self._to_latent_stats(z_mean) z_norm = (z_mean - mean) / std @@ -193,7 +191,7 @@ class Cosmos25RMSNorm(nn.Module): def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None: super().__init__() broadcastable_dims = (1, 1, 1) if not images else (1, 1) - shape = (dim, *broadcastable_dims) if channel_first else (dim,) + shape = (dim, *broadcastable_dims) if channel_first else (dim, ) self.channel_first = channel_first self.scale = dim**0.5 @@ -254,7 +252,8 @@ class Cosmos25Resample(nn.Module): else: cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": - cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2) if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) @@ -287,6 +286,7 @@ class Cosmos25Resample(nn.Module): class Cosmos25ResidualBlock(nn.Module): + def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.0) -> None: super().__init__() self.in_dim = in_dim @@ -309,7 +309,8 @@ class Cosmos25ResidualBlock(nn.Module): idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 @@ -334,13 +335,7 @@ class Cosmos25AttentionBlock(nn.Module): b, c, t, h, w = x.size() x2 = rearrange(x, "b c t h w -> (b t) c h w") x2 = self.norm(x2) - q, k, v = ( - self.to_qkv(x2) - .reshape(b * t, 1, c * 3, -1) - .permute(0, 1, 3, 2) - .contiguous() - .chunk(3, dim=-1) - ) + q, k, v = (self.to_qkv(x2).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1)) x2 = F.scaled_dot_product_attention(q, k, v) x2 = x2.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) x2 = self.proj(x2) @@ -349,6 +344,7 @@ class Cosmos25AttentionBlock(nn.Module): class Cosmos25Encoder3d(nn.Module): + def __init__( self, dim: int = 96, @@ -417,7 +413,8 @@ class Cosmos25Encoder3d(nn.Module): idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 @@ -427,6 +424,7 @@ class Cosmos25Encoder3d(nn.Module): class Cosmos25Decoder3d(nn.Module): + def __init__( self, dim: int = 96, @@ -439,7 +437,7 @@ class Cosmos25Decoder3d(nn.Module): ) -> None: super().__init__() dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] - scale = 1.0 / 2 ** (len(dim_mult) - 2) + scale = 1.0 / 2**(len(dim_mult) - 2) self.conv1 = Cosmos25CausalConv3d(z_dim, dims[0], 3, padding=1) self.middle = nn.Sequential( @@ -495,7 +493,8 @@ class Cosmos25Decoder3d(nn.Module): idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 @@ -655,7 +654,7 @@ class Cosmos25WanVAE(nn.Module): out = self._i0_encode(x) else: out_ = self.encoder( - x[:, :, 1 + self.temporal_window * (i - 1) : 1 + self.temporal_window * i, :, :], + x[:, :, 1 + self.temporal_window * (i - 1):1 + self.temporal_window * i, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx, ) @@ -664,7 +663,7 @@ class Cosmos25WanVAE(nn.Module): if (t - 1) % self.temporal_window: self._enc_conv_idx = [0] out_ = self.encoder( - x[:, :, 1 + self.temporal_window * (iters - 1) :, :, :], + x[:, :, 1 + self.temporal_window * (iters - 1):, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx, ) @@ -691,7 +690,7 @@ class Cosmos25WanVAE(nn.Module): if i == 0: out = self._i0_decode(x) else: - out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) + out_ = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) out = torch.cat([out, out_], 2) self.clear_cache() @@ -715,12 +714,3 @@ class Cosmos25WanVAE(nn.Module): @property def latent_ch(self) -> int: return 16 - - - - - - - - - diff --git a/fastvideo/models/vaes/flux2_components.py b/fastvideo/models/vaes/flux2_components.py index 0e174ba54..de427542a 100644 --- a/fastvideo/models/vaes/flux2_components.py +++ b/fastvideo/models/vaes/flux2_components.py @@ -18,7 +18,7 @@ class AutoencoderKLOutput: latent_dist: DiagonalGaussianDistribution def __getitem__(self, idx: int): - return (self.latent_dist,)[idx] + return (self.latent_dist, )[idx] def __getattr__(self, name: str): # Existing local Flux2 parity tests used `vae.encode(x).mean` while @@ -48,7 +48,11 @@ def get_activation(act_fn: str) -> nn.Module: class AttnProcessor: - def __call__(self, attn: "Attention", hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + + def __call__(self, + attn: "Attention", + hidden_states: torch.Tensor, + temb: Optional[torch.Tensor] = None) -> torch.Tensor: return attn._forward(hidden_states, temb=temb) @@ -61,6 +65,7 @@ CROSS_ATTENTION_PROCESSORS = frozenset({AttnProcessor}) class Attention(nn.Module): + def __init__( self, query_dim: int, @@ -91,11 +96,8 @@ class Attention(nn.Module): self.spatial_norm = None if spatial_norm_dim is not None: raise ValueError("Flux2 VAE does not use spatial attention norm in this port") - self.group_norm = ( - nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True) - if norm_num_groups is not None - else None - ) + self.group_norm = (nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True) + if norm_num_groups is not None else None) self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias) self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias) self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias) @@ -142,7 +144,13 @@ class Attention(nn.Module): class Downsample2D(nn.Module): - def __init__(self, channels: int, use_conv: bool = False, out_channels: Optional[int] = None, padding: int = 1, name: str = "conv"): + + def __init__(self, + channels: int, + use_conv: bool = False, + out_channels: Optional[int] = None, + padding: int = 1, + name: str = "conv"): super().__init__() self.channels = channels self.out_channels = out_channels or channels @@ -170,6 +178,7 @@ class Downsample2D(nn.Module): class Upsample2D(nn.Module): + def __init__(self, channels: int, use_conv: bool = False, out_channels: Optional[int] = None, name: str = "conv"): super().__init__() self.channels = channels @@ -201,6 +210,7 @@ class Upsample2D(nn.Module): class ResnetBlock2D(nn.Module): + def __init__( self, *, @@ -232,7 +242,8 @@ class ResnetBlock2D(nn.Module): self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) if temb_channels is not None: - self.time_emb_proj = nn.Linear(temb_channels, out_channels if time_embedding_norm == "default" else 2 * out_channels) + self.time_emb_proj = nn.Linear(temb_channels, + out_channels if time_embedding_norm == "default" else 2 * out_channels) else: self.time_emb_proj = None self.norm2 = nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True) @@ -243,7 +254,12 @@ class ResnetBlock2D(nn.Module): self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut self.conv_shortcut = None if self.use_in_shortcut: - self.conv_shortcut = nn.Conv2d(in_channels, conv_2d_out_channels, kernel_size=1, stride=1, padding=0, bias=conv_shortcut_bias) + self.conv_shortcut = nn.Conv2d(in_channels, + conv_2d_out_channels, + kernel_size=1, + stride=1, + padding=0, + bias=conv_shortcut_bias) def forward(self, input_tensor: torch.Tensor, temb: Optional[torch.Tensor] = None, *args, **kwargs) -> torch.Tensor: hidden_states = self.norm1(input_tensor) @@ -271,6 +287,7 @@ class ResnetBlock2D(nn.Module): class UNetMidBlock2D(nn.Module): + def __init__( self, in_channels: int, @@ -324,10 +341,7 @@ class UNetMidBlock2D(nn.Module): bias=True, upcast_softmax=True, _from_deprecated_attn_block=True, - ) - if self.add_attention - else None - ) + ) if self.add_attention else None) resnets.append( ResnetBlock2D( in_channels=in_channels, @@ -340,8 +354,7 @@ class UNetMidBlock2D(nn.Module): non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, - ) - ) + )) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) self.gradient_checkpointing = False @@ -356,7 +369,18 @@ class UNetMidBlock2D(nn.Module): class DownEncoderBlock2D(nn.Module): - def __init__(self, in_channels: int, out_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, resnet_act_fn: str = "swish", resnet_groups: int = 32, add_downsample: bool = True, downsample_padding: int = 1, **_: object): + + def __init__(self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + add_downsample: bool = True, + downsample_padding: int = 1, + **_: object): super().__init__() self.resnets = nn.ModuleList([ ResnetBlock2D( @@ -367,10 +391,11 @@ class DownEncoderBlock2D(nn.Module): groups=resnet_groups, dropout=dropout, non_linearity=resnet_act_fn, - ) - for i in range(num_layers) + ) for i in range(num_layers) ]) - self.downsamplers = nn.ModuleList([Downsample2D(out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op")]) if add_downsample else None + self.downsamplers = nn.ModuleList([ + Downsample2D(out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op") + ]) if add_downsample else None def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: for resnet in self.resnets: @@ -382,6 +407,7 @@ class DownEncoderBlock2D(nn.Module): class AttnDownEncoderBlock2D(DownEncoderBlock2D): + def __init__(self, in_channels: int, out_channels: int, attention_head_dim: int = 1, **kwargs: object): super().__init__(in_channels=in_channels, out_channels=out_channels, **kwargs) resnet_groups = int(kwargs.get("resnet_groups", 32)) @@ -390,8 +416,16 @@ class AttnDownEncoderBlock2D(DownEncoderBlock2D): if attention_head_dim is None: attention_head_dim = out_channels self.attentions = nn.ModuleList([ - Attention(out_channels, heads=out_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor, eps=resnet_eps, norm_num_groups=resnet_groups, residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True) - for _ in self.resnets + Attention(out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True) for _ in self.resnets ]) def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: @@ -405,7 +439,18 @@ class AttnDownEncoderBlock2D(DownEncoderBlock2D): class UpDecoderBlock2D(nn.Module): - def __init__(self, in_channels: int, out_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, resnet_act_fn: str = "swish", resnet_groups: int = 32, add_upsample: bool = True, temb_channels: Optional[int] = None, **_: object): + + def __init__(self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + add_upsample: bool = True, + temb_channels: Optional[int] = None, + **_: object): super().__init__() self.resnets = nn.ModuleList([ ResnetBlock2D( @@ -416,10 +461,10 @@ class UpDecoderBlock2D(nn.Module): groups=resnet_groups, dropout=dropout, non_linearity=resnet_act_fn, - ) - for i in range(num_layers) + ) for i in range(num_layers) ]) - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) if add_upsample else None + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels) + ]) if add_upsample else None self.resolution_idx = None def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: @@ -432,6 +477,7 @@ class UpDecoderBlock2D(nn.Module): class AttnUpDecoderBlock2D(UpDecoderBlock2D): + def __init__(self, in_channels: int, out_channels: int, attention_head_dim: int = 1, **kwargs: object): super().__init__(in_channels=in_channels, out_channels=out_channels, **kwargs) resnet_groups = int(kwargs.get("resnet_groups", 32)) @@ -440,8 +486,16 @@ class AttnUpDecoderBlock2D(UpDecoderBlock2D): if attention_head_dim is None: attention_head_dim = out_channels self.attentions = nn.ModuleList([ - Attention(out_channels, heads=out_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor, eps=resnet_eps, norm_num_groups=resnet_groups, residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True) - for _ in self.resnets + Attention(out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True) for _ in self.resnets ]) def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: @@ -473,7 +527,17 @@ def get_up_block(up_block_type: str, **kwargs: object) -> nn.Module: class Encoder(nn.Module): - def __init__(self, in_channels: int = 3, out_channels: int = 3, down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",), block_out_channels: Tuple[int, ...] = (64,), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", double_z: bool = True, mid_block_add_attention: bool = True): + + def __init__(self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D", ), + block_out_channels: Tuple[int, ...] = (64, ), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + double_z: bool = True, + mid_block_add_attention: bool = True): super().__init__() self.layers_per_block = layers_per_block self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) @@ -483,8 +547,27 @@ class Encoder(nn.Module): input_channel = output_channel output_channel = block_out_channels[i] is_final_block = i == len(block_out_channels) - 1 - self.down_blocks.append(get_down_block(down_block_type, num_layers=self.layers_per_block, in_channels=input_channel, out_channels=output_channel, add_downsample=not is_final_block, resnet_eps=1e-6, downsample_padding=0, resnet_act_fn=act_fn, resnet_groups=norm_num_groups, attention_head_dim=output_channel, temb_channels=None)) - self.mid_block = UNetMidBlock2D(in_channels=block_out_channels[-1], resnet_eps=1e-6, resnet_act_fn=act_fn, output_scale_factor=1, resnet_time_scale_shift="default", attention_head_dim=block_out_channels[-1], resnet_groups=norm_num_groups, temb_channels=None, add_attention=mid_block_add_attention) + self.down_blocks.append( + get_down_block(down_block_type, + num_layers=self.layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=not is_final_block, + resnet_eps=1e-6, + downsample_padding=0, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + attention_head_dim=output_channel, + temb_channels=None)) + self.mid_block = UNetMidBlock2D(in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default", + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=None, + add_attention=mid_block_add_attention) self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) self.conv_act = nn.SiLU() conv_out_channels = 2 * out_channels if double_z else out_channels @@ -502,21 +585,51 @@ class Encoder(nn.Module): class Decoder(nn.Module): - def __init__(self, in_channels: int = 3, out_channels: int = 3, up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), block_out_channels: Tuple[int, ...] = (64,), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", norm_type: str = "group", mid_block_add_attention: bool = True): + + def __init__(self, + in_channels: int = 3, + out_channels: int = 3, + up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D", ), + block_out_channels: Tuple[int, ...] = (64, ), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + norm_type: str = "group", + mid_block_add_attention: bool = True): super().__init__() if norm_type != "group": raise ValueError("Flux2 VAE Decoder only supports group norm in this port") self.layers_per_block = layers_per_block self.conv_in = nn.Conv2d(in_channels, block_out_channels[-1], kernel_size=3, stride=1, padding=1) self.up_blocks = nn.ModuleList([]) - self.mid_block = UNetMidBlock2D(in_channels=block_out_channels[-1], resnet_eps=1e-6, resnet_act_fn=act_fn, output_scale_factor=1, resnet_time_scale_shift="default", attention_head_dim=block_out_channels[-1], resnet_groups=norm_num_groups, temb_channels=None, add_attention=mid_block_add_attention) + self.mid_block = UNetMidBlock2D(in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default", + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=None, + add_attention=mid_block_add_attention) reversed_block_out_channels = list(reversed(block_out_channels)) output_channel = reversed_block_out_channels[0] for i, up_block_type in enumerate(up_block_types): prev_output_channel = output_channel output_channel = reversed_block_out_channels[i] is_final_block = i == len(block_out_channels) - 1 - self.up_blocks.append(get_up_block(up_block_type, num_layers=self.layers_per_block + 1, in_channels=prev_output_channel, out_channels=output_channel, prev_output_channel=prev_output_channel, add_upsample=not is_final_block, resnet_eps=1e-6, resnet_act_fn=act_fn, resnet_groups=norm_num_groups, attention_head_dim=output_channel, temb_channels=None, resnet_time_scale_shift=norm_type)) + self.up_blocks.append( + get_up_block(up_block_type, + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + prev_output_channel=prev_output_channel, + add_upsample=not is_final_block, + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + attention_head_dim=output_channel, + temb_channels=None, + resnet_time_scale_shift=norm_type)) self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) self.conv_act = nn.SiLU() self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1) diff --git a/fastvideo/models/vaes/flux2vae.py b/fastvideo/models/vaes/flux2vae.py index fce5ede88..c1325deec 100644 --- a/fastvideo/models/vaes/flux2vae.py +++ b/fastvideo/models/vaes/flux2vae.py @@ -63,7 +63,7 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): batch_norm_eps: float = arch_config.batch_norm_eps batch_norm_momentum: float = arch_config.batch_norm_momentum patch_size: Tuple[int, int] = arch_config.patch_size - + # pass init params to Encoder self.encoder = Encoder( in_channels=in_channels, @@ -89,16 +89,8 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): mid_block_add_attention=mid_block_add_attention, ) - self.quant_conv = ( - nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) - if use_quant_conv - else None - ) - self.post_quant_conv = ( - nn.Conv2d(latent_channels, latent_channels, 1) - if use_post_quant_conv - else None - ) + self.quant_conv = (nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) if use_quant_conv else None) + self.post_quant_conv = (nn.Conv2d(latent_channels, latent_channels, 1) if use_post_quant_conv else None) self.bn = nn.BatchNorm2d( math.prod(patch_size) * latent_channels, @@ -113,14 +105,8 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): # only relevant if vae tiling is enabled self.tile_sample_min_size = sample_size - sample_size_val = ( - sample_size[0] - if isinstance(sample_size, (list, tuple)) - else sample_size - ) - self.tile_latent_min_size = int( - sample_size_val / (2 ** (len(block_out_channels) - 1)) - ) + sample_size_val = (sample_size[0] if isinstance(sample_size, (list, tuple)) else sample_size) + self.tile_latent_min_size = int(sample_size_val / (2**(len(block_out_channels) - 1))) self.tile_overlap_factor = 0.25 @property @@ -153,9 +139,7 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): return processors # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]] - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -173,8 +157,7 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): if isinstance(processor, dict) and len(processor) != count: raise ValueError( f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) + f" number of attention layers: {count}. Please make sure to pass {count} processor classes.") def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): @@ -194,15 +177,9 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): """ Disables custom attention processors and sets the default attention implementation. """ - if all( - proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS - for proc in self.attn_processors.values() - ): + if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): processor = AttnAddedKVProcessor() - elif all( - proc.__class__ in CROSS_ATTENTION_PROCESSORS - for proc in self.attn_processors.values() - ): + elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): processor = AttnProcessor() else: raise ValueError( @@ -214,9 +191,7 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): def _encode(self, x: torch.Tensor) -> torch.Tensor: batch_size, num_channels, height, width = x.shape - if self.use_tiling and ( - width > self.tile_sample_min_size or height > self.tile_sample_min_size - ): + if self.use_tiling and (width > self.tile_sample_min_size or height > self.tile_sample_min_size): return self._tiled_encode(x) enc = self.encoder(x) @@ -225,9 +200,9 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): return enc - def encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + def encode(self, + x: torch.Tensor, + return_dict: bool = True) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: """ Encode a batch of images into latents. @@ -253,17 +228,12 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): posterior = DiagonalGaussianDistribution(h) if not return_dict: - return (posterior,) + return (posterior, ) return AutoencoderKLOutput(latent_dist=posterior) - def _decode( - self, z: torch.Tensor, return_dict: bool = True - ) -> Union[DecoderOutput, torch.Tensor]: - if self.use_tiling and ( - z.shape[-1] > self.tile_latent_min_size - or z.shape[-2] > self.tile_latent_min_size - ): + def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size): return self.tiled_decode(z, return_dict=return_dict) if self.post_quant_conv is not None: @@ -272,13 +242,14 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): dec = self.decoder(z) if not return_dict: - return (dec,) + return (dec, ) return DecoderOutput(sample=dec) - def decode( - self, z: torch.FloatTensor, return_dict: bool = True, generator=None - ) -> Union[DecoderOutput, torch.FloatTensor]: + def decode(self, + z: torch.FloatTensor, + return_dict: bool = True, + generator=None) -> Union[DecoderOutput, torch.FloatTensor]: """ Decode a batch of images. @@ -300,28 +271,20 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): decoded = self._decode(z).sample if not return_dict: - return (decoded,) + return (decoded, ) return DecoderOutput(sample=decoded) - def blend_v( - self, a: torch.Tensor, b: torch.Tensor, blend_extent: int - ) -> torch.Tensor: + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: blend_extent = min(a.shape[2], b.shape[2], blend_extent) for y in range(blend_extent): - b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[ - :, :, y, : - ] * (y / blend_extent) + b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent) return b - def blend_h( - self, a: torch.Tensor, b: torch.Tensor, blend_extent: int - ) -> torch.Tensor: + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: blend_extent = min(a.shape[3], b.shape[3], blend_extent) for x in range(blend_extent): - b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[ - :, :, :, x - ] * (x / blend_extent) + b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent) return b def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor: @@ -353,8 +316,8 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): tile = x[ :, :, - i : i + self.tile_sample_min_size, - j : j + self.tile_sample_min_size, + i:i + self.tile_sample_min_size, + j:j + self.tile_sample_min_size, ] tile = self.encoder(tile) if self.quant_conv is not None: @@ -377,9 +340,7 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): enc = torch.cat(result_rows, dim=2) return enc - def tiled_encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> AutoencoderKLOutput: + def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutput: r"""Encode a batch of images using a tiled encoder. When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several @@ -416,8 +377,8 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): tile = x[ :, :, - i : i + self.tile_sample_min_size, - j : j + self.tile_sample_min_size, + i:i + self.tile_sample_min_size, + j:j + self.tile_sample_min_size, ] tile = self.encoder(tile) if self.quant_conv is not None: @@ -441,13 +402,11 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): posterior = DiagonalGaussianDistribution(moments) if not return_dict: - return (posterior,) + return (posterior, ) return AutoencoderKLOutput(latent_dist=posterior) - def tiled_decode( - self, z: torch.Tensor, return_dict: bool = True - ) -> Union[DecoderOutput, torch.Tensor]: + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: r""" Decode a batch of images using a tiled decoder. @@ -474,8 +433,8 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): tile = z[ :, :, - i : i + self.tile_latent_min_size, - j : j + self.tile_latent_min_size, + i:i + self.tile_latent_min_size, + j:j + self.tile_latent_min_size, ] if self.post_quant_conv is not None: tile = self.post_quant_conv(tile) @@ -497,7 +456,7 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): dec = torch.cat(result_rows, dim=2) if not return_dict: - return (dec,) + return (dec, ) return DecoderOutput(sample=dec) @@ -525,7 +484,7 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): dec = self.decode(z).sample if not return_dict: - return (dec,) + return (dec, ) return DecoderOutput(sample=dec) diff --git a/fastvideo/models/vaes/gamecraftvae.py b/fastvideo/models/vaes/gamecraftvae.py index 815f09b9e..e687c346c 100644 --- a/fastvideo/models/vaes/gamecraftvae.py +++ b/fastvideo/models/vaes/gamecraftvae.py @@ -42,7 +42,7 @@ class EncoderCausal3D(nn.Module): self, in_channels: int = 3, out_channels: int = 16, - down_block_types: Tuple[str, ...] = ("DownEncoderBlockCausal3D",), + down_block_types: Tuple[str, ...] = ("DownEncoderBlockCausal3D", ), block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), layers_per_block: int = 2, norm_num_groups: int = 32, @@ -57,9 +57,11 @@ class EncoderCausal3D(nn.Module): super().__init__() self.layers_per_block = layers_per_block - self.conv_in = CausalConv3d( - in_channels, block_out_channels[0], kernel_size=3, stride=1, disable_causal=disable_causal - ) + self.conv_in = CausalConv3d(in_channels, + block_out_channels[0], + kernel_size=3, + stride=1, + disable_causal=disable_causal) self.down_blocks = nn.ModuleList([]) output_channel = block_out_channels[0] @@ -72,9 +74,7 @@ class EncoderCausal3D(nn.Module): if time_compression_ratio == 4: add_spatial = bool(i < num_spatial) - add_time = bool( - i >= (len(block_out_channels) - 1 - num_time) and not is_final_block - ) + add_time = bool(i >= (len(block_out_channels) - 1 - num_time) and not is_final_block) elif time_compression_ratio == 8: add_spatial = bool(i < num_spatial) add_time = bool(i < num_time) @@ -82,7 +82,7 @@ class EncoderCausal3D(nn.Module): raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}") downsample_stride_HW = (2, 2) if add_spatial else (1, 1) - downsample_stride_T = (2,) if add_time else (1,) + downsample_stride_T = (2, ) if add_time else (1, ) downsample_stride = tuple(downsample_stride_T + downsample_stride_HW) down_block = DownEncoderBlockCausal3D( @@ -112,9 +112,7 @@ class EncoderCausal3D(nn.Module): causal_attention=mid_block_causal_attn, ) - self.conv_norm_out = nn.GroupNorm( - num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6 - ) + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) self.conv_act = nn.SiLU() conv_out_channels = 2 * out_channels if double_z else out_channels self.conv_out = CausalConv3d( @@ -142,7 +140,7 @@ class DecoderCausal3D(nn.Module): self, in_channels: int = 16, out_channels: int = 3, - up_block_types: Tuple[str, ...] = ("UpDecoderBlockCausal3D",), + up_block_types: Tuple[str, ...] = ("UpDecoderBlockCausal3D", ), block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), layers_per_block: int = 2, norm_num_groups: int = 32, @@ -189,9 +187,7 @@ class DecoderCausal3D(nn.Module): if time_compression_ratio == 4: add_spatial = bool(i < num_spatial) - add_time = bool( - i >= len(block_out_channels) - 1 - num_time and not is_final_block - ) + add_time = bool(i >= len(block_out_channels) - 1 - num_time and not is_final_block) elif time_compression_ratio == 8: add_spatial = bool(i >= len(block_out_channels) - num_spatial) add_time = bool(i >= len(block_out_channels) - num_time) @@ -199,7 +195,7 @@ class DecoderCausal3D(nn.Module): raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}") upsample_HW = (2, 2) if add_spatial else (1, 1) - upsample_T = (2,) if add_time else (1,) + upsample_T = (2, ) if add_time else (1, ) upsample_scale_factor = tuple(upsample_T + upsample_HW) up_block = UpDecoderBlockCausal3D( @@ -217,13 +213,9 @@ class DecoderCausal3D(nn.Module): self.up_blocks.append(up_block) output_channel = output_channel - self.conv_norm_out = nn.GroupNorm( - num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6 - ) + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) self.conv_act = nn.SiLU() - self.conv_out = CausalConv3d( - block_out_channels[0], out_channels, kernel_size=3, disable_causal=disable_causal - ) + self.conv_out = CausalConv3d(block_out_channels[0], out_channels, kernel_size=3, disable_causal=disable_causal) def forward( self, @@ -283,27 +275,23 @@ class GameCraftVAE(nn.Module): mid_block_causal_attn=getattr(arch, "mid_block_causal_attn", False), ) - self.quant_conv = nn.Conv3d( - 2 * arch.latent_channels, 2 * arch.latent_channels, kernel_size=1 - ) - self.post_quant_conv = nn.Conv3d( - arch.latent_channels, arch.latent_channels, kernel_size=1 - ) - + self.quant_conv = nn.Conv3d(2 * arch.latent_channels, 2 * arch.latent_channels, kernel_size=1) + self.post_quant_conv = nn.Conv3d(arch.latent_channels, arch.latent_channels, kernel_size=1) + # Scaling factor for latent normalization (required for decoding stage) self.scaling_factor = arch.scaling_factor - + # Tiling support - matches official GameCraft VAE settings self._tiling_enabled = False self.tile_overlap_factor = 0.25 - + # Temporal tiling params (for >64 output frames) self.tile_sample_min_tsize = 64 # Minimum sample temporal size (video frames) self.tile_latent_min_tsize = 16 # = 64 // 4 (time_compression_ratio) - + # Spatial tiling params - use small tiles to reduce memory self.tile_sample_min_size = 256 # Minimum spatial tile size in pixel space - self.tile_latent_min_size = 32 # = 256 // 8 (spatial_compression_ratio) + self.tile_latent_min_size = 32 # = 256 // 8 (spatial_compression_ratio) def encode(self, x: torch.Tensor) -> AutoencoderKLOutput: """Encode to latent distribution.""" @@ -329,7 +317,7 @@ class GameCraftVAE(nn.Module): # Check if spatial tiling needed (large H or W) if z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size: return self._spatial_tiled_decode(z) - + z = self.post_quant_conv(z) dec = self.decoder(z, latent_embeds=None) return dec @@ -341,21 +329,21 @@ class GameCraftVAE(nn.Module): Only used when T > tile_latent_min_tsize (16). """ B, C, T, H, W = z.shape - + # Use the pre-configured tiling parameters overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor)) blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor) t_limit = self.tile_sample_min_tsize - blend_extent - + row = [] for i in range(0, T, overlap_size): - tile = z[:, :, i : i + self.tile_latent_min_tsize + 1, :, :] + tile = z[:, :, i:i + self.tile_latent_min_tsize + 1, :, :] tile = self.post_quant_conv(tile) decoded = self.decoder(tile, latent_embeds=None) if i > 0: decoded = decoded[:, :, 1:, :, :] # Skip first frame for non-first tiles row.append(decoded) - + # Blend overlapping regions result_row = [] for i, tile in enumerate(row): @@ -363,10 +351,10 @@ class GameCraftVAE(nn.Module): tile = self._blend_t(row[i - 1], tile, blend_extent) result_row.append(tile[:, :, :t_limit, :, :]) else: - result_row.append(tile[:, :, :t_limit+1, :, :]) - + result_row.append(tile[:, :, :t_limit + 1, :, :]) + return torch.cat(result_row, dim=2) - + def _spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor: """Decode latents in spatial tiles with overlapping and blending. @@ -375,18 +363,18 @@ class GameCraftVAE(nn.Module): overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) row_limit = self.tile_sample_min_size - blend_extent - + # Split z into overlapping tiles and decode them separately rows = [] for i in range(0, z.shape[-2], overlap_size): row = [] for j in range(0, z.shape[-1], overlap_size): - tile = z[:, :, :, i: i + self.tile_latent_min_size, j: j + self.tile_latent_min_size] + tile = z[:, :, :, i:i + self.tile_latent_min_size, j:j + self.tile_latent_min_size] tile = self.post_quant_conv(tile) decoded = self.decoder(tile, latent_embeds=None) row.append(decoded) rows.append(row) - + # Blend overlapping regions result_rows = [] for i, row in enumerate(rows): @@ -399,53 +387,53 @@ class GameCraftVAE(nn.Module): tile = self._blend_h(row[j - 1], tile, blend_extent) result_row.append(tile[:, :, :, :row_limit, :row_limit]) result_rows.append(torch.cat(result_row, dim=-1)) - + return torch.cat(result_rows, dim=-2) - + def _blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: """Blend two tensors along temporal dimension.""" blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) if blend_extent == 0: return b - + a_region = a[..., -blend_extent:, :, :] b_region = b[..., :blend_extent, :, :] - + weights = torch.arange(blend_extent, device=a.device, dtype=a.dtype) / blend_extent weights = weights.view(1, 1, blend_extent, 1, 1) - + blended = a_region * (1 - weights) + b_region * weights b[..., :blend_extent, :, :] = blended return b - + def _blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: """Blend two tensors along vertical (height) dimension.""" blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) if blend_extent == 0: return b - + a_region = a[..., -blend_extent:, :] b_region = b[..., :blend_extent, :] - + weights = torch.arange(blend_extent, device=a.device, dtype=a.dtype) / blend_extent weights = weights.view(1, 1, 1, blend_extent, 1) - + blended = a_region * (1 - weights) + b_region * weights b[..., :blend_extent, :] = blended return b - + def _blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: """Blend two tensors along horizontal (width) dimension.""" blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) if blend_extent == 0: return b - + a_region = a[..., -blend_extent:] b_region = b[..., :blend_extent] - + weights = torch.arange(blend_extent, device=a.device, dtype=a.dtype) / blend_extent weights = weights.view(1, 1, 1, 1, blend_extent) - + blended = a_region * (1 - weights) + b_region * weights b[..., :blend_extent] = blended return b diff --git a/fastvideo/models/vaes/gamecraftvae_blocks.py b/fastvideo/models/vaes/gamecraftvae_blocks.py index 7ce9bb7d7..7f6bedd7a 100644 --- a/fastvideo/models/vaes/gamecraftvae_blocks.py +++ b/fastvideo/models/vaes/gamecraftvae_blocks.py @@ -12,14 +12,12 @@ import torch.nn.functional as F from einops import rearrange -def prepare_causal_attention_mask( - n_frame: int, n_hw: int, dtype, device, batch_size: Optional[int] = None -): +def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: Optional[int] = None): seq_len = n_frame * n_hw mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device) for i in range(seq_len): i_frame = i // n_hw - mask[i, : (i_frame + 1) * n_hw] = 0 + mask[i, :(i_frame + 1) * n_hw] = 0 if batch_size is not None: mask = mask.unsqueeze(0).expand(batch_size, -1, -1) return mask @@ -47,9 +45,7 @@ class CausalConv3d(nn.Module): else: padding = (k // 2, k // 2, k // 2, k // 2, k - 1, 0) self.time_causal_padding = padding - self.conv = nn.Conv3d( - chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs - ) + self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) def forward(self, x: torch.Tensor) -> torch.Tensor: x = F.pad(x, self.time_causal_padding, mode=self.pad_mode) @@ -89,13 +85,13 @@ class UpsampleCausal3D(nn.Module): """Causal 3D upsampling - matches official (has .conv when use_conv=True).""" def __init__( - self, - channels: int, - out_channels: Optional[int] = None, - kernel_size: int = 3, - upsample_factor: Tuple[int, int, int] = (2, 2, 2), - disable_causal: bool = False, - bias: bool = True, + self, + channels: int, + out_channels: Optional[int] = None, + kernel_size: int = 3, + upsample_factor: Tuple[int, int, int] = (2, 2, 2), + disable_causal: bool = False, + bias: bool = True, ): super().__init__() self.out_channels = out_channels or channels @@ -123,17 +119,12 @@ class UpsampleCausal3D(nn.Module): if not self.disable_causal and T > 1: first_h, other_h = hidden_states.split((1, T - 1), dim=2) - other_h = F.interpolate( - other_h, scale_factor=self.upsample_factor, mode="nearest" - ) - first_h = F.interpolate( - first_h.squeeze(2), scale_factor=self.upsample_factor[1:], mode="nearest" - ).unsqueeze(2) + other_h = F.interpolate(other_h, scale_factor=self.upsample_factor, mode="nearest") + first_h = F.interpolate(first_h.squeeze(2), scale_factor=self.upsample_factor[1:], + mode="nearest").unsqueeze(2) hidden_states = torch.cat((first_h, other_h), dim=2) else: - hidden_states = F.interpolate( - hidden_states, scale_factor=self.upsample_factor, mode="nearest" - ) + hidden_states = F.interpolate(hidden_states, scale_factor=self.upsample_factor, mode="nearest") if dtype == torch.bfloat16: hidden_states = hidden_states.to(dtype) @@ -156,9 +147,7 @@ class GameCraftVAEAttention(nn.Module): self.heads = heads self.dim_head = dim_head inner_dim = heads * dim_head - self.group_norm = nn.GroupNorm( - norm_num_groups or in_channels, in_channels, eps=eps - ) + self.group_norm = nn.GroupNorm(norm_num_groups or in_channels, in_channels, eps=eps) self.to_q = nn.Linear(in_channels, inner_dim, bias=bias) self.to_k = nn.Linear(in_channels, inner_dim, bias=bias) self.to_v = nn.Linear(in_channels, inner_dim, bias=bias) @@ -172,9 +161,7 @@ class GameCraftVAEAttention(nn.Module): residual = hidden_states batch_size, seq_len, _ = hidden_states.shape - hidden_states = self.group_norm( - hidden_states.permute(0, 2, 1) - ).permute(0, 2, 1) + hidden_states = self.group_norm(hidden_states.permute(0, 2, 1)).permute(0, 2, 1) q = self.to_q(hidden_states) k = self.to_k(hidden_states) @@ -217,11 +204,8 @@ class ResnetBlockCausal3D(nn.Module): self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps) self.conv2 = CausalConv3d(out_channels, out_channels, 3, 1, disable_causal=disable_causal) self.dropout = nn.Dropout(dropout) - self.conv_shortcut = ( - CausalConv3d(in_channels, out_channels, 1, 1, disable_causal=disable_causal) - if in_channels != out_channels - else None - ) + self.conv_shortcut = (CausalConv3d(in_channels, out_channels, 1, 1, disable_causal=disable_causal) + if in_channels != out_channels else None) self.nonlinearity = getattr(F, non_linearity, F.silu) def forward( @@ -275,8 +259,7 @@ class UNetMidBlockCausal3D(nn.Module): dropout=0.0, non_linearity=resnet_act_fn, disable_causal=disable_causal, - ) - ) + )) for _ in range(num_layers): if add_attention: @@ -288,8 +271,7 @@ class UNetMidBlockCausal3D(nn.Module): eps=resnet_eps, norm_num_groups=resnet_groups, bias=True, - ) - ) + )) else: self.attentions.append(None) self.resnets.append( @@ -302,8 +284,7 @@ class UNetMidBlockCausal3D(nn.Module): dropout=0.0, non_linearity=resnet_act_fn, disable_causal=disable_causal, - ) - ) + )) def forward( self, @@ -316,15 +297,15 @@ class UNetMidBlockCausal3D(nn.Module): B, C, T, H, W = hidden_states.shape hidden_states = rearrange(hidden_states, "b c f h w -> b (f h w) c") if self.causal_attention: - mask = prepare_causal_attention_mask( - T, H * W, hidden_states.dtype, hidden_states.device, batch_size=B - ) + mask = prepare_causal_attention_mask(T, + H * W, + hidden_states.dtype, + hidden_states.device, + batch_size=B) else: mask = None hidden_states = attn(hidden_states, attention_mask=mask) - hidden_states = rearrange( - hidden_states, "b (f h w) c -> b c f h w", f=T, h=H, w=W - ) + hidden_states = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=T, h=H, w=W) hidden_states = resnet(hidden_states, temb) return hidden_states @@ -359,8 +340,7 @@ class DownEncoderBlockCausal3D(nn.Module): dropout=0.0, non_linearity=resnet_act_fn, disable_causal=disable_causal, - ) - ) + )) self.downsamplers = None if add_downsample: @@ -374,9 +354,7 @@ class DownEncoderBlockCausal3D(nn.Module): ) ]) - def forward( - self, hidden_states: torch.Tensor, scale: float = 1.0 - ) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, scale: float = 1.0) -> torch.Tensor: for resnet in self.resnets: hidden_states = resnet(hidden_states, temb=None, scale=scale) if self.downsamplers is not None: @@ -415,8 +393,7 @@ class UpDecoderBlockCausal3D(nn.Module): dropout=0.0, non_linearity=resnet_act_fn, disable_causal=disable_causal, - ) - ) + )) self.upsamplers = None if add_upsample: diff --git a/fastvideo/models/vaes/gen3c_tokenizer_vae.py b/fastvideo/models/vaes/gen3c_tokenizer_vae.py index 3c633286a..4057f1658 100644 --- a/fastvideo/models/vaes/gen3c_tokenizer_vae.py +++ b/fastvideo/models/vaes/gen3c_tokenizer_vae.py @@ -61,10 +61,8 @@ class _JITGen3CTokenizerInner(nn.Module): super().__init__() self._dtype = dtype self._forced_bf16 = False - self.encoder = torch.jit.load(encoder_path, map_location=device).eval().to( - device=device, dtype=dtype) - self.decoder = torch.jit.load(decoder_path, map_location=device).eval().to( - device=device, dtype=dtype) + self.encoder = torch.jit.load(encoder_path, map_location=device).eval().to(device=device, dtype=dtype) + self.decoder = torch.jit.load(decoder_path, map_location=device).eval().to(device=device, dtype=dtype) latent_mean, latent_std = torch.load(mean_std_path, map_location="cpu") latent_mean = latent_mean.view(latent_channels, -1)[:, :latent_chunk_duration] @@ -72,14 +70,12 @@ class _JITGen3CTokenizerInner(nn.Module): self.register_buffer( "_latent_mean", - latent_mean.to(torch.float32).view(1, latent_channels, - latent_chunk_duration, 1, 1), + latent_mean.to(torch.float32).view(1, latent_channels, latent_chunk_duration, 1, 1), persistent=False, ) self.register_buffer( "_latent_std", - latent_std.to(torch.float32).view(1, latent_channels, - latent_chunk_duration, 1, 1), + latent_std.to(torch.float32).view(1, latent_channels, latent_chunk_duration, 1, 1), persistent=False, ) @@ -92,10 +88,8 @@ class _JITGen3CTokenizerInner(nn.Module): if t < mean.shape[2]: return mean[:, :, :t], std[:, :, :t] # fallback for non-default lengths - mean = torch.nn.functional.interpolate( - mean, size=(t, 1, 1), mode="trilinear", align_corners=False) - std = torch.nn.functional.interpolate( - std, size=(t, 1, 1), mode="trilinear", align_corners=False) + mean = torch.nn.functional.interpolate(mean, size=(t, 1, 1), mode="trilinear", align_corners=False) + std = torch.nn.functional.interpolate(std, size=(t, 1, 1), mode="trilinear", align_corners=False) return mean, std @staticmethod @@ -113,9 +107,7 @@ class _JITGen3CTokenizerInner(nn.Module): self.decoder = self.decoder.to(dtype=torch.bfloat16) self._dtype = torch.bfloat16 self._forced_bf16 = True - logger.warning( - "GEN3C tokenizer JIT hit fp16/bf16 mismatch; coercing tokenizer encoder/decoder to bf16." - ) + logger.warning("GEN3C tokenizer JIT hit fp16/bf16 mismatch; coercing tokenizer encoder/decoder to bf16.") def encode(self, x: torch.Tensor) -> _TensorLatentDist: enc_dtype, enc_device = self._module_dtype_device(self.encoder) @@ -225,12 +217,9 @@ class AutoencoderKLGen3CTokenizer(nn.Module): if self._target_temporal_compression == 2 * self._inner_temporal_compression: b, c, t, h, w = z_target.shape t_inner = 2 * t - 1 - out = torch.empty( - b, c, t_inner, h, w, device=z_target.device, dtype=z_target.dtype) + out = torch.empty(b, c, t_inner, h, w, device=z_target.device, dtype=z_target.dtype) out[:, :, 0::2, :, :] = z_target - out[:, :, 1::2, :, :] = 0.5 * ( - z_target[:, :, :-1, :, :] + z_target[:, :, 1:, :, :] - ) + out[:, :, 1::2, :, :] = 0.5 * (z_target[:, :, :-1, :, :] + z_target[:, :, 1:, :, :]) return out.contiguous() # Generic fallback: linear index interpolation in time. @@ -354,8 +343,7 @@ class AutoencoderKLGen3CTokenizer(nn.Module): dtype=dtype, device=device, latent_channels=16, - latent_chunk_duration=1 + (pixel_chunk_duration - 1) // - target_temporal_compression, + latent_chunk_duration=1 + (pixel_chunk_duration - 1) // target_temporal_compression, ) return cls( inner, diff --git a/fastvideo/models/vaes/hunyuan15vae.py b/fastvideo/models/vaes/hunyuan15vae.py index 40085b6bb..fb14e1750 100644 --- a/fastvideo/models/vaes/hunyuan15vae.py +++ b/fastvideo/models/vaes/hunyuan15vae.py @@ -27,7 +27,9 @@ from fastvideo.layers.activation import get_act_fn from fastvideo.configs.models.vaes import Hunyuan15VAEConfig from fastvideo.models.vaes.common import ParallelTiledVAE + class HunyuanVideo15CausalConv3d(nn.Module): + def __init__( self, in_channels: int, @@ -75,7 +77,7 @@ class HunyuanVideo15RMS_norm(nn.Module): def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None: super().__init__() broadcastable_dims = (1, 1, 1) if not images else (1, 1) - shape = (dim, *broadcastable_dims) if channel_first else (dim,) + shape = (dim, *broadcastable_dims) if channel_first else (dim, ) self.channel_first = channel_first self.scale = dim**0.5 @@ -87,6 +89,7 @@ class HunyuanVideo15RMS_norm(nn.Module): class HunyuanVideo15AttnBlock(nn.Module): + def __init__(self, in_channels: int): super().__init__() self.in_channels = in_channels @@ -116,7 +119,7 @@ class HunyuanVideo15AttnBlock(nn.Module): mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device) for i in range(seq_len): i_frame = i // n_hw - mask[i, : (i_frame + 1) * n_hw] = 0 + mask[i, :(i_frame + 1) * n_hw] = 0 if batch_size is not None: mask = mask.unsqueeze(0).expand(batch_size, -1, -1) return mask @@ -136,9 +139,11 @@ class HunyuanVideo15AttnBlock(nn.Module): key = key.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous() value = value.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous() - attention_mask = self.prepare_causal_attention_mask( - frames, height * width, query.dtype, query.device, batch_size=batch_size - ) + attention_mask = self.prepare_causal_attention_mask(frames, + height * width, + query.dtype, + query.device, + batch_size=batch_size) x = nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask) @@ -151,6 +156,7 @@ class HunyuanVideo15AttnBlock(nn.Module): class HunyuanVideo15Upsample(nn.Module): + def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: bool = True): super().__init__() factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2 @@ -184,7 +190,7 @@ class HunyuanVideo15Upsample(nn.Module): if self.add_temporal_upsample: h_first = h[:, :, :1, :, :] h_first = self._dcae_upsample_rearrange(h_first, r1=1, r2=2, r3=2) - h_first = h_first[:, : h_first.shape[1] // 2] + h_first = h_first[:, :h_first.shape[1] // 2] h_next = h[:, :, 1:, :, :] h_next = self._dcae_upsample_rearrange(h_next, r1=r1, r2=2, r3=2) h = torch.cat([h_first, h_next], dim=2) @@ -207,6 +213,7 @@ class HunyuanVideo15Upsample(nn.Module): class HunyuanVideo15Downsample(nn.Module): + def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample: bool = True): super().__init__() factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2 @@ -260,6 +267,7 @@ class HunyuanVideo15Downsample(nn.Module): class HunyuanVideo15ResnetBlock(nn.Module): + def __init__( self, in_channels: int, @@ -299,6 +307,7 @@ class HunyuanVideo15ResnetBlock(nn.Module): class HunyuanVideo15MidBlock(nn.Module): + def __init__( self, in_channels: int, @@ -309,12 +318,10 @@ class HunyuanVideo15MidBlock(nn.Module): self.add_attention = add_attention # There is always at least one resnet - resnets = [ - HunyuanVideo15ResnetBlock( - in_channels=in_channels, - out_channels=in_channels, - ) - ] + resnets = [HunyuanVideo15ResnetBlock( + in_channels=in_channels, + out_channels=in_channels, + )] attentions = [] for _ in range(num_layers): @@ -323,12 +330,10 @@ class HunyuanVideo15MidBlock(nn.Module): else: attentions.append(None) - resnets.append( - HunyuanVideo15ResnetBlock( - in_channels=in_channels, - out_channels=in_channels, - ) - ) + resnets.append(HunyuanVideo15ResnetBlock( + in_channels=in_channels, + out_channels=in_channels, + )) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) @@ -347,6 +352,7 @@ class HunyuanVideo15MidBlock(nn.Module): class HunyuanVideo15DownBlock3D(nn.Module): + def __init__( self, in_channels: int, @@ -360,25 +366,21 @@ class HunyuanVideo15DownBlock3D(nn.Module): for i in range(num_layers): in_channels = in_channels if i == 0 else out_channels - resnets.append( - HunyuanVideo15ResnetBlock( - in_channels=in_channels, - out_channels=out_channels, - ) - ) + resnets.append(HunyuanVideo15ResnetBlock( + in_channels=in_channels, + out_channels=out_channels, + )) self.resnets = nn.ModuleList(resnets) if downsample_out_channels is not None: - self.downsamplers = nn.ModuleList( - [ - HunyuanVideo15Downsample( - out_channels, - out_channels=downsample_out_channels, - add_temporal_downsample=add_temporal_downsample, - ) - ] - ) + self.downsamplers = nn.ModuleList([ + HunyuanVideo15Downsample( + out_channels, + out_channels=downsample_out_channels, + add_temporal_downsample=add_temporal_downsample, + ) + ]) else: self.downsamplers = None @@ -396,6 +398,7 @@ class HunyuanVideo15DownBlock3D(nn.Module): class HunyuanVideo15UpBlock3D(nn.Module): + def __init__( self, in_channels: int, @@ -410,25 +413,21 @@ class HunyuanVideo15UpBlock3D(nn.Module): for i in range(num_layers): input_channels = in_channels if i == 0 else out_channels - resnets.append( - HunyuanVideo15ResnetBlock( - in_channels=input_channels, - out_channels=out_channels, - ) - ) + resnets.append(HunyuanVideo15ResnetBlock( + in_channels=input_channels, + out_channels=out_channels, + )) self.resnets = nn.ModuleList(resnets) if upsample_out_channels is not None: - self.upsamplers = nn.ModuleList( - [ - HunyuanVideo15Upsample( - out_channels, - out_channels=upsample_out_channels, - add_temporal_upsample=add_temporal_upsample, - ) - ] - ) + self.upsamplers = nn.ModuleList([ + HunyuanVideo15Upsample( + out_channels, + out_channels=upsample_out_channels, + add_temporal_upsample=add_temporal_upsample, + ) + ]) else: self.upsamplers = None @@ -668,7 +667,7 @@ class AutoencoderKLHunyuanVideo15(nn.Module, ParallelTiledVAE): # The minimal tile height and width for spatial tiling to be used self.tile_sample_min_height = 256 self.tile_sample_min_width = 256 - self.tile_sample_min_num_frames = 2000 # Fill in a random large number, as hy1.5 vae does not use temporal tiling + self.tile_sample_min_num_frames = 2000 # Fill in a random large number, as hy1.5 vae does not use temporal tiling def _encode(self, x: torch.Tensor) -> torch.Tensor: x = self.encoder(x) @@ -699,5 +698,6 @@ class AutoencoderKLHunyuanVideo15(nn.Module, ParallelTiledVAE): dec = self.decode(z) return dec + # Entry point for model registry EntryClass = AutoencoderKLHunyuanVideo15 diff --git a/fastvideo/models/vaes/hunyuanvae.py b/fastvideo/models/vaes/hunyuanvae.py index b3d04b487..f6520c794 100644 --- a/fastvideo/models/vaes/hunyuanvae.py +++ b/fastvideo/models/vaes/hunyuanvae.py @@ -25,12 +25,11 @@ from fastvideo.layers.activation import get_act_fn from fastvideo.models.vaes.common import ParallelTiledVAE -def prepare_causal_attention_mask( - num_frames: int, - height_width: int, - dtype: torch.dtype, - device: torch.device, - batch_size: int | None = None) -> torch.Tensor: +def prepare_causal_attention_mask(num_frames: int, + height_width: int, + dtype: torch.dtype, + device: torch.device, + batch_size: int | None = None) -> torch.Tensor: indices = torch.arange(1, num_frames + 1, dtype=torch.int32, device=device) indices_blocks = indices.repeat_interleave(height_width) x, y = torch.meshgrid(indices_blocks, indices_blocks, indexing="xy") @@ -43,8 +42,7 @@ def prepare_causal_attention_mask( class HunyuanVAEAttention(nn.Module): - def __init__(self, in_channels, heads, dim_head, eps, norm_num_groups, - bias) -> None: + def __init__(self, in_channels, heads, dim_head, eps, norm_num_groups, bias) -> None: super().__init__() self.in_channels = in_channels self.heads = heads @@ -59,25 +57,17 @@ class HunyuanVAEAttention(nn.Module): self.to_q = nn.Linear(in_channels, inner_dim, bias=bias) self.to_k = nn.Linear(in_channels, inner_dim, bias=bias) self.to_v = nn.Linear(in_channels, inner_dim, bias=bias) - self.to_out = nn.Sequential(nn.Linear(inner_dim, in_channels, - bias=bias)) + self.to_out = nn.Sequential(nn.Linear(inner_dim, in_channels, bias=bias)) # Optional normalization layers - self.group_norm = nn.GroupNorm(norm_num_groups, - in_channels, - eps=eps, - affine=True) + self.group_norm = nn.GroupNorm(norm_num_groups, in_channels, eps=eps, affine=True) - def forward(self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor: residual = hidden_states batch_size, sequence_length, _ = hidden_states.shape - hidden_states = self.group_norm(hidden_states.transpose(1, - 2)).transpose( - 1, 2) + hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) # Project to query, key, value query = self.to_q(hidden_states) @@ -100,8 +90,7 @@ class HunyuanVAEAttention(nn.Module): is_causal=False) # Reshape back - hidden_states = hidden_states.transpose(1, 2).reshape( - batch_size, -1, self.heads * head_dim) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) hidden_states = hidden_states.to(query.dtype) # Linear projection @@ -128,8 +117,7 @@ class HunyuanVideoCausalConv3d(nn.Module): ) -> None: super().__init__() - kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance( - kernel_size, int) else kernel_size + kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size self.pad_mode = pad_mode self.time_causal_padding = ( @@ -141,18 +129,10 @@ class HunyuanVideoCausalConv3d(nn.Module): 0, ) - self.conv = nn.Conv3d(in_channels, - out_channels, - kernel_size, - stride, - padding, - dilation, - bias=bias) + self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = F.pad(hidden_states, - self.time_causal_padding, - mode=self.pad_mode) + hidden_states = F.pad(hidden_states, self.time_causal_padding, mode=self.pad_mode) return self.conv(hidden_states) @@ -172,19 +152,13 @@ class HunyuanVideoUpsampleCausal3D(nn.Module): out_channels = out_channels or in_channels self.upsample_factor = upsample_factor - self.conv = HunyuanVideoCausalConv3d(in_channels, - out_channels, - kernel_size, - stride, - bias=bias) + self.conv = HunyuanVideoCausalConv3d(in_channels, out_channels, kernel_size, stride, bias=bias) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_frames = hidden_states.size(2) - first_frame, other_frames = hidden_states.split((1, num_frames - 1), - dim=2) - first_frame = F.interpolate(first_frame.squeeze(2), - scale_factor=self.upsample_factor[1:], + first_frame, other_frames = hidden_states.split((1, num_frames - 1), dim=2) + first_frame = F.interpolate(first_frame.squeeze(2), scale_factor=self.upsample_factor[1:], mode="nearest").unsqueeze(2) if num_frames > 1: @@ -195,9 +169,7 @@ class HunyuanVideoUpsampleCausal3D(nn.Module): # `vae.enable_tiling()` first. If that doesn't work, open an issue at: # https://github.com/huggingface/diffusers/issues other_frames = other_frames.contiguous() - other_frames = F.interpolate(other_frames, - scale_factor=self.upsample_factor, - mode="nearest") + other_frames = F.interpolate(other_frames, scale_factor=self.upsample_factor, mode="nearest") hidden_states = torch.cat((first_frame, other_frames), dim=2) else: hidden_states = first_frame @@ -220,12 +192,7 @@ class HunyuanVideoDownsampleCausal3D(nn.Module): super().__init__() out_channels = out_channels or channels - self.conv = HunyuanVideoCausalConv3d(channels, - out_channels, - kernel_size, - stride, - padding, - bias=bias) + self.conv = HunyuanVideoCausalConv3d(channels, out_channels, kernel_size, stride, padding, bias=bias) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.conv(hidden_states) @@ -249,18 +216,15 @@ class HunyuanVideoResnetBlockCausal3D(nn.Module): self.nonlinearity = get_act_fn(non_linearity) self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps, affine=True) - self.conv1 = HunyuanVideoCausalConv3d(in_channels, out_channels, 3, 1, - 0) + self.conv1 = HunyuanVideoCausalConv3d(in_channels, out_channels, 3, 1, 0) self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps, affine=True) self.dropout = nn.Dropout(dropout) - self.conv2 = HunyuanVideoCausalConv3d(out_channels, out_channels, 3, 1, - 0) + self.conv2 = HunyuanVideoCausalConv3d(out_channels, out_channels, 3, 1, 0) self.conv_shortcut = None if in_channels != out_channels: - self.conv_shortcut = HunyuanVideoCausalConv3d( - in_channels, out_channels, 1, 1, 0) + self.conv_shortcut = HunyuanVideoCausalConv3d(in_channels, out_channels, 1, 1, 0) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = hidden_states.contiguous() @@ -296,8 +260,7 @@ class HunyuanVideoMidBlock3D(nn.Module): attention_head_dim: int = 1, ) -> None: super().__init__() - resnet_groups = resnet_groups if resnet_groups is not None else min( - in_channels // 4, 32) + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) self.add_attention = add_attention # There is always at least one resnet @@ -344,50 +307,36 @@ class HunyuanVideoMidBlock3D(nn.Module): def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - self.resnets[0], hidden_states) + hidden_states = self._gradient_checkpointing_func(self.resnets[0], hidden_states) - for attn, resnet in zip(self.attentions, - self.resnets[1:], - strict=True): + for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True): if attn is not None: batch_size, num_channels, num_frames, height, width = hidden_states.shape - hidden_states = hidden_states.permute(0, 2, 3, 4, - 1).flatten(1, 3) - attention_mask = prepare_causal_attention_mask( - num_frames, - height * width, - hidden_states.dtype, - hidden_states.device, - batch_size=batch_size) - hidden_states = attn(hidden_states, - attention_mask=attention_mask) - hidden_states = hidden_states.unflatten( - 1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) + hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3) + attention_mask = prepare_causal_attention_mask(num_frames, + height * width, + hidden_states.dtype, + hidden_states.device, + batch_size=batch_size) + hidden_states = attn(hidden_states, attention_mask=attention_mask) + hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) - hidden_states = self._gradient_checkpointing_func( - resnet, hidden_states) + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) else: hidden_states = self.resnets[0](hidden_states) - for attn, resnet in zip(self.attentions, - self.resnets[1:], - strict=True): + for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True): if attn is not None: batch_size, num_channels, num_frames, height, width = hidden_states.shape - hidden_states = hidden_states.permute(0, 2, 3, 4, - 1).flatten(1, 3) - attention_mask = prepare_causal_attention_mask( - num_frames, - height * width, - hidden_states.dtype, - hidden_states.device, - batch_size=batch_size) - hidden_states = attn(hidden_states, - attention_mask=attention_mask) - hidden_states = hidden_states.unflatten( - 1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) + hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3) + attention_mask = prepare_causal_attention_mask(num_frames, + height * width, + hidden_states.dtype, + hidden_states.device, + batch_size=batch_size) + hidden_states = attn(hidden_states, attention_mask=attention_mask) + hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) hidden_states = resnet(hidden_states) @@ -443,8 +392,7 @@ class HunyuanVideoDownBlock3D(nn.Module): def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if torch.is_grad_enabled() and self.gradient_checkpointing: for resnet in self.resnets: - hidden_states = self._gradient_checkpointing_func( - resnet, hidden_states) + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) else: for resnet in self.resnets: hidden_states = resnet(hidden_states) @@ -504,8 +452,7 @@ class HunyuanVideoUpBlock3D(nn.Module): def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if torch.is_grad_enabled() and self.gradient_checkpointing: for resnet in self.resnets: - hidden_states = self._gradient_checkpointing_func( - resnet, hidden_states) + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) else: for resnet in self.resnets: @@ -544,51 +491,40 @@ class HunyuanVideoEncoder3D(nn.Module): ) -> None: super().__init__() - self.conv_in = HunyuanVideoCausalConv3d(in_channels, - block_out_channels[0], - kernel_size=3, - stride=1) + self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1) self.mid_block: HunyuanVideoMidBlock3D | None = None self.down_blocks = nn.ModuleList([]) output_channel = block_out_channels[0] for i, down_block_type in enumerate(down_block_types): if down_block_type != "HunyuanVideoDownBlock3D": - raise ValueError( - f"Unsupported down_block_type: {down_block_type}") + raise ValueError(f"Unsupported down_block_type: {down_block_type}") input_channel = output_channel output_channel = block_out_channels[i] is_final_block = i == len(block_out_channels) - 1 - num_spatial_downsample_layers = int( - np.log2(spatial_compression_ratio)) - num_time_downsample_layers = int( - np.log2(temporal_compression_ratio)) + num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_downsample_layers = int(np.log2(temporal_compression_ratio)) if temporal_compression_ratio == 4: add_spatial_downsample = bool(i < num_spatial_downsample_layers) - add_time_downsample = bool(i >= (len(block_out_channels) - 1 - - num_time_downsample_layers) + add_time_downsample = bool(i >= (len(block_out_channels) - 1 - num_time_downsample_layers) and not is_final_block) elif temporal_compression_ratio == 8: add_spatial_downsample = bool(i < num_spatial_downsample_layers) add_time_downsample = bool(i < num_time_downsample_layers) else: - raise ValueError( - f"Unsupported time_compression_ratio: {temporal_compression_ratio}" - ) + raise ValueError(f"Unsupported time_compression_ratio: {temporal_compression_ratio}") downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) downsample_stride_T = (2, ) if add_time_downsample else (1, ) - downsample_stride = tuple(downsample_stride_T + - downsample_stride_HW) + downsample_stride = tuple(downsample_stride_T + downsample_stride_HW) down_block = HunyuanVideoDownBlock3D( num_layers=layers_per_block, in_channels=input_channel, out_channels=output_channel, - add_downsample=bool(add_spatial_downsample - or add_time_downsample), + add_downsample=bool(add_spatial_downsample or add_time_downsample), resnet_eps=1e-6, resnet_act_fn=act_fn, resnet_groups=norm_num_groups, @@ -607,15 +543,11 @@ class HunyuanVideoEncoder3D(nn.Module): add_attention=mid_block_add_attention, ) - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], - num_groups=norm_num_groups, - eps=1e-6) + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) self.conv_act = nn.SiLU() conv_out_channels = 2 * out_channels if double_z else out_channels - self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[-1], - conv_out_channels, - kernel_size=3) + self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3) self.gradient_checkpointing = False @@ -624,11 +556,9 @@ class HunyuanVideoEncoder3D(nn.Module): if torch.is_grad_enabled() and self.gradient_checkpointing: for down_block in self.down_blocks: - hidden_states = self._gradient_checkpointing_func( - down_block, hidden_states) + hidden_states = self._gradient_checkpointing_func(down_block, hidden_states) - hidden_states = self._gradient_checkpointing_func( - self.mid_block, hidden_states) + hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) else: for down_block in self.down_blocks: hidden_states = down_block(hidden_states) @@ -668,10 +598,7 @@ class HunyuanVideoDecoder3D(nn.Module): super().__init__() self.layers_per_block = layers_per_block - self.conv_in = HunyuanVideoCausalConv3d(in_channels, - block_out_channels[-1], - kernel_size=3, - stride=1) + self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1) self.up_blocks = nn.ModuleList([]) # mid @@ -694,25 +621,19 @@ class HunyuanVideoDecoder3D(nn.Module): prev_output_channel = output_channel output_channel = reversed_block_out_channels[i] is_final_block = i == len(block_out_channels) - 1 - num_spatial_upsample_layers = int( - np.log2(spatial_compression_ratio)) + num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio)) num_time_upsample_layers = int(np.log2(time_compression_ratio)) if time_compression_ratio == 4: add_spatial_upsample = bool(i < num_spatial_upsample_layers) - add_time_upsample = bool( - i >= len(block_out_channels) - 1 - num_time_upsample_layers - and not is_final_block) + add_time_upsample = bool(i >= len(block_out_channels) - 1 - num_time_upsample_layers + and not is_final_block) else: - raise ValueError( - f"Unsupported time_compression_ratio: {time_compression_ratio}" - ) + raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}") - upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, - 1) + upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1) upsample_scale_factor_T = (2, ) if add_time_upsample else (1, ) - upsample_scale_factor = tuple(upsample_scale_factor_T + - upsample_scale_factor_HW) + upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW) up_block = HunyuanVideoUpBlock3D( num_layers=self.layers_per_block + 1, @@ -729,13 +650,9 @@ class HunyuanVideoDecoder3D(nn.Module): prev_output_channel = output_channel # out - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], - num_groups=norm_num_groups, - eps=1e-6) + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) self.conv_act = nn.SiLU() - self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[0], - out_channels, - kernel_size=3) + self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[0], out_channels, kernel_size=3) self.gradient_checkpointing = False @@ -743,12 +660,10 @@ class HunyuanVideoDecoder3D(nn.Module): hidden_states = self.conv_in(hidden_states) if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - self.mid_block, hidden_states) + hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) for up_block in self.up_blocks: - hidden_states = self._gradient_checkpointing_func( - up_block, hidden_states) + hidden_states = self._gradient_checkpointing_func(up_block, hidden_states) else: hidden_states = self.mid_block(hidden_states) @@ -799,9 +714,7 @@ class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE): temporal_compression_ratio=config.temporal_compression_ratio, spatial_compression_ratio=config.spatial_compression_ratio, ) - self.quant_conv = nn.Conv3d(2 * config.latent_channels, - 2 * config.latent_channels, - kernel_size=1) + self.quant_conv = nn.Conv3d(2 * config.latent_channels, 2 * config.latent_channels, kernel_size=1) if config.load_decoder: self.decoder = HunyuanVideoDecoder3D( @@ -816,9 +729,7 @@ class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE): spatial_compression_ratio=config.spatial_compression_ratio, mid_block_add_attention=config.mid_block_add_attention, ) - self.post_quant_conv = nn.Conv3d(config.latent_channels, - config.latent_channels, - kernel_size=1) + self.post_quant_conv = nn.Conv3d(config.latent_channels, config.latent_channels, kernel_size=1) def _encode(self, x: torch.Tensor) -> torch.Tensor: x = self.encoder(x) @@ -848,5 +759,6 @@ class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE): dec = self.decode(z) return dec + # Entry point for model registry EntryClass = AutoencoderKLHunyuanVideo diff --git a/fastvideo/models/vaes/hyworldvae.py b/fastvideo/models/vaes/hyworldvae.py index 62e9ea7b3..4e95225bf 100644 --- a/fastvideo/models/vaes/hyworldvae.py +++ b/fastvideo/models/vaes/hyworldvae.py @@ -35,9 +35,10 @@ from fastvideo.models.vaes.hunyuan15vae import ( # Cache size for temporal feature caching (number of frames to cache) CACHE_T = 2 + class HYWorldCausalConv3d(nn.Module): """Causal Conv3d with optional cache support for temporal feature caching.""" - + def __init__( self, in_channels: int, @@ -60,8 +61,8 @@ class HYWorldCausalConv3d(nn.Module): kernel_size[0] // 2, # W_right (spatial) kernel_size[1] // 2, # H_left (spatial) kernel_size[1] // 2, # H_right (spatial) - kernel_size[2] - 1, # T_left (temporal causal padding) - 0, # T_right (no future padding for causal) + kernel_size[2] - 1, # T_left (temporal causal padding) + 0, # T_right (no future padding for causal) ) self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias) @@ -76,21 +77,21 @@ class HYWorldCausalConv3d(nn.Module): When provided, uses cached frames instead of padding for temporal dimension. """ padding = list(self.time_causal_padding) - + if cache_x is not None and self.time_causal_padding[4] > 0: # Has temporal padding and cache cache_x = cache_x.to(hidden_states.device) # Concatenate cached frames with current input on temporal dimension hidden_states = torch.cat([cache_x, hidden_states], dim=2) # Reduce temporal padding since we have cached frames padding[4] -= cache_x.shape[2] - + hidden_states = F.pad(hidden_states, padding, mode=self.pad_mode) return self.conv(hidden_states) class HYWorldUpsample(nn.Module): """Hierarchical upsampling with temporal/spatial support and optional caching.""" - + def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: bool = True): super().__init__() factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2 @@ -99,7 +100,7 @@ class HYWorldUpsample(nn.Module): self.repeats = factor * out_channels // in_channels def forward( - self, + self, x: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None, @@ -133,7 +134,7 @@ class HYWorldUpsample(nn.Module): if first_chunk: # First chunk: only spatial upsample h = rearrange(h, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) - h = h[:, : h.shape[1] // 2] + h = h[:, :h.shape[1] // 2] # Compute the shortcut part shortcut = rearrange(x, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) shortcut = shortcut.repeat_interleave(repeats=self.repeats // 2, dim=1) @@ -146,7 +147,7 @@ class HYWorldUpsample(nn.Module): # First frame: only spatial upsample h_first = rearrange(h_first, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) - h_first = h_first[:, : h_first.shape[1] // 2] + h_first = h_first[:, :h_first.shape[1] // 2] shortcut_first = rearrange(x_first, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) shortcut_first = shortcut_first.repeat_interleave(repeats=self.repeats // 2, dim=1) out_first = h_first + shortcut_first @@ -167,13 +168,13 @@ class HYWorldUpsample(nn.Module): h = rearrange(h, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) shortcut = x.repeat_interleave(repeats=self.repeats, dim=1) shortcut = rearrange(shortcut, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) - + return h + shortcut class HYWorldDownsample(nn.Module): """Hierarchical downsampling with temporal/spatial support and optional caching.""" - + def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample: bool = True): super().__init__() factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2 @@ -183,7 +184,7 @@ class HYWorldDownsample(nn.Module): self.group_size = factor * in_channels // out_channels def forward( - self, + self, x: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None, @@ -197,7 +198,7 @@ class HYWorldDownsample(nn.Module): feat_idx: List containing current cache index [idx] """ r1 = 2 if self.add_temporal_downsample else 1 - + # Apply conv with caching if feat_cache is not None and feat_idx is not None: idx = feat_idx[0] @@ -211,8 +212,8 @@ class HYWorldDownsample(nn.Module): feat_cache[idx] = cache_x feat_idx[0] += 1 else: - h = self.conv(x) # Change the channel, ready for spatial or temporal downsample - + h = self.conv(x) # Change the channel, ready for spatial or temporal downsample + if self.add_temporal_downsample: if x.shape[2] == 1: # Single frame: only spatial downsample @@ -239,6 +240,7 @@ class HYWorldDownsample(nn.Module): class HYWorldResnetBlock(nn.Module): + def __init__( self, in_channels: int, @@ -259,9 +261,9 @@ class HYWorldResnetBlock(nn.Module): self.conv2 = HYWorldCausalConv3d(out_channels, out_channels, kernel_size=3) if in_channels != out_channels: self.nin_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) - + def forward( - self, + self, hidden_states: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None, @@ -278,7 +280,7 @@ class HYWorldResnetBlock(nn.Module): hidden_states = self.norm1(hidden_states) hidden_states = self.nonlinearity(hidden_states) - + # apply the feature cacheing mechanism if feat_cache is not None and feat_idx is not None: # Retrieve the current layer index. @@ -286,7 +288,7 @@ class HYWorldResnetBlock(nn.Module): # Clone the last CACHE_T frames from the current input to store for the next step. cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() - + # Handle boundary conditions: if the current temporal chunk is too short (< 2 frames) # and we have a previous cache, prepend the last frame of the previous cache. # This ensures sufficient temporal context for the convolution kernel. @@ -296,10 +298,10 @@ class HYWorldResnetBlock(nn.Module): [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2, ) - + # Apply the convolution layer using `hidden_states`and the *previous* cached state. hidden_states = self.conv1(hidden_states, feat_cache[idx]) - + # Update the cache for this layer with the newly prepared context ('cache_x'). feat_cache[idx] = cache_x @@ -334,7 +336,7 @@ class HYWorldResnetBlock(nn.Module): class HYWorldMidBlock(nn.Module): """Mid block with attention and resnet blocks, with optional caching support.""" - + def __init__( self, in_channels: int, @@ -345,12 +347,10 @@ class HYWorldMidBlock(nn.Module): self.add_attention = add_attention # There is always at least one resnet - resnets = [ - HYWorldResnetBlock( - in_channels=in_channels, - out_channels=in_channels, - ) - ] + resnets = [HYWorldResnetBlock( + in_channels=in_channels, + out_channels=in_channels, + )] attentions = [] for _ in range(num_layers): @@ -359,12 +359,10 @@ class HYWorldMidBlock(nn.Module): else: attentions.append(None) - resnets.append( - HYWorldResnetBlock( - in_channels=in_channels, - out_channels=in_channels, - ) - ) + resnets.append(HYWorldResnetBlock( + in_channels=in_channels, + out_channels=in_channels, + )) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) @@ -372,7 +370,7 @@ class HYWorldMidBlock(nn.Module): self.gradient_checkpointing = False def forward( - self, + self, hidden_states: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None, @@ -397,7 +395,7 @@ class HYWorldMidBlock(nn.Module): class HYWorldDownBlock3D(nn.Module): """Down block with resnet blocks and optional downsampling, with caching support.""" - + def __init__( self, in_channels: int, @@ -411,32 +409,28 @@ class HYWorldDownBlock3D(nn.Module): for i in range(num_layers): in_channels = in_channels if i == 0 else out_channels - resnets.append( - HYWorldResnetBlock( - in_channels=in_channels, - out_channels=out_channels, - ) - ) + resnets.append(HYWorldResnetBlock( + in_channels=in_channels, + out_channels=out_channels, + )) self.resnets = nn.ModuleList(resnets) if downsample_out_channels is not None: - self.downsamplers = nn.ModuleList( - [ - HYWorldDownsample( - out_channels, - out_channels=downsample_out_channels, - add_temporal_downsample=add_temporal_downsample, - ) - ] - ) + self.downsamplers = nn.ModuleList([ + HYWorldDownsample( + out_channels, + out_channels=downsample_out_channels, + add_temporal_downsample=add_temporal_downsample, + ) + ]) else: self.downsamplers = None self.gradient_checkpointing = False def forward( - self, + self, hidden_states: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None, @@ -461,7 +455,7 @@ class HYWorldDownBlock3D(nn.Module): class HYWorldUpBlock3D(nn.Module): """Up block with resnet blocks and optional upsampling, with caching support.""" - + def __init__( self, in_channels: int, @@ -476,32 +470,28 @@ class HYWorldUpBlock3D(nn.Module): for i in range(num_layers): input_channels = in_channels if i == 0 else out_channels - resnets.append( - HYWorldResnetBlock( - in_channels=input_channels, - out_channels=out_channels, - ) - ) + resnets.append(HYWorldResnetBlock( + in_channels=input_channels, + out_channels=out_channels, + )) self.resnets = nn.ModuleList(resnets) if upsample_out_channels is not None: - self.upsamplers = nn.ModuleList( - [ - HYWorldUpsample( - out_channels, - out_channels=upsample_out_channels, - add_temporal_upsample=add_temporal_upsample, - ) - ] - ) + self.upsamplers = nn.ModuleList([ + HYWorldUpsample( + out_channels, + out_channels=upsample_out_channels, + add_temporal_upsample=add_temporal_upsample, + ) + ]) else: self.upsamplers = None self.gradient_checkpointing = False def forward( - self, + self, hidden_states: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None, @@ -592,7 +582,7 @@ class HYWorldEncoder3D(nn.Module): self.gradient_checkpointing = False def forward( - self, + self, hidden_states: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None, @@ -720,7 +710,7 @@ class HYWorldDecoder3D(nn.Module): self.gradient_checkpointing = False def forward( - self, + self, hidden_states: torch.Tensor, feat_cache: Optional[List[Optional[torch.Tensor]]] = None, feat_idx: Optional[List[int]] = None, @@ -745,8 +735,7 @@ class HYWorldDecoder3D(nn.Module): dim=2, ) hidden_states = self.conv_in(hidden_states, feat_cache[idx]) + hidden_states.repeat_interleave( - repeats=self.repeat, dim=1 - ) + repeats=self.repeat, dim=1) feat_cache[idx] = cache_x feat_idx[0] += 1 else: @@ -779,7 +768,7 @@ class HYWorldDecoder3D(nn.Module): feat_idx[0] += 1 else: hidden_states = self.conv_out(hidden_states) - + return hidden_states @@ -792,8 +781,8 @@ class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): _supports_gradient_checkpointing = True def __init__( - self, - config: Hunyuan15VAEConfig, # HYWorld use the same VAE architecture as HunyuanVideo-1.5 + self, + config: Hunyuan15VAEConfig, # HYWorld use the same VAE architecture as HunyuanVideo-1.5 ) -> None: nn.Module.__init__(self) ParallelTiledVAE.__init__(self, config) @@ -829,7 +818,7 @@ class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): # The minimal tile height and width for spatial tiling to be used self.tile_sample_min_height = 256 self.tile_sample_min_width = 256 - self.tile_sample_min_num_frames = 2000 # Fill in a random large number, as hy1.5 vae does not use temporal tiling + self.tile_sample_min_num_frames = 2000 # Fill in a random large number, as hy1.5 vae does not use temporal tiling # Cache-related attributes (initialized in clear_cache) self._conv_num: int = 0 @@ -841,16 +830,10 @@ class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): # Precompute and cache conv counts for encoder and decoder for clear_cache speedup self._cached_conv_counts = { - "decoder": ( - sum(1 for m in self.decoder.modules() if isinstance(m, HYWorldCausalConv3d)) - if self.decoder is not None - else 0 - ), - "encoder": ( - sum(1 for m in self.encoder.modules() if isinstance(m, HYWorldCausalConv3d)) - if self.encoder is not None - else 0 - ), + "decoder": (sum(1 for m in self.decoder.modules() + if isinstance(m, HYWorldCausalConv3d)) if self.decoder is not None else 0), + "encoder": (sum(1 for m in self.encoder.modules() + if isinstance(m, HYWorldCausalConv3d)) if self.encoder is not None else 0), } def clear_cache(self) -> None: @@ -864,7 +847,7 @@ class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): self._conv_num = self._cached_conv_counts["decoder"] self._conv_idx = [0] self._feat_map: List[Optional[torch.Tensor]] = [None] * self._conv_num - + # Cache for encoder self._enc_conv_num = self._cached_conv_counts["encoder"] self._enc_conv_idx = [0] @@ -888,7 +871,7 @@ class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): _, _, num_frame, _, _ = x.shape self.clear_cache() - + # Process in chunks: first frame alone, then groups of 4 frames iter_ = 1 + (num_frame - 1) // 4 for i in range(iter_): @@ -903,7 +886,7 @@ class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): else: # Subsequent frames in groups of 4 out_ = self.encoder( - x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx, ) @@ -930,14 +913,14 @@ class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): _, _, num_frame, _, _ = z.shape self.clear_cache() - + # Process one frame at a time with caching for i in range(num_frame): self._conv_idx = [0] if i == 0: # First frame with first_chunk=True out = self.decoder( - z[:, :, i : i + 1, :, :], + z[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=True, @@ -945,7 +928,7 @@ class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): else: # Subsequent frames out_ = self.decoder( - z[:, :, i : i + 1, :, :], + z[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=False, @@ -984,5 +967,6 @@ class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): return dec + # Entry point for model registry EntryClass = AutoencoderKLHYWorld diff --git a/fastvideo/models/vaes/lingbotworld2_wanvae.py b/fastvideo/models/vaes/lingbotworld2_wanvae.py index 0c8e9485c..979fd2410 100644 --- a/fastvideo/models/vaes/lingbotworld2_wanvae.py +++ b/fastvideo/models/vaes/lingbotworld2_wanvae.py @@ -22,8 +22,7 @@ class CausalConv3d(nn.Conv3d): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self._padding = (self.padding[2], self.padding[2], self.padding[1], - self.padding[1], 2 * self.padding[0], 0) + self._padding = (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0) self.padding = (0, 0, 0) def forward(self, x, cache_x=None): @@ -42,7 +41,7 @@ class RMS_norm(nn.Module): def __init__(self, dim, channel_first=True, images=True, bias=False): super().__init__() broadcastable_dims = (1, 1, 1) if not images else (1, 1) - shape = (dim, *broadcastable_dims) if channel_first else (dim,) + shape = (dim, *broadcastable_dims) if channel_first else (dim, ) self.channel_first = channel_first self.scale = dim**0.5 @@ -50,9 +49,7 @@ class RMS_norm(nn.Module): self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. def forward(self, x): - return F.normalize( - x, dim=(1 if self.channel_first else - -1)) * self.scale * self.gamma + self.bias + return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias class Upsample(nn.Upsample): @@ -67,34 +64,25 @@ class Upsample(nn.Upsample): class Resample(nn.Module): def __init__(self, dim, mode): - assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d', - 'downsample3d') + assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d', 'downsample3d') super().__init__() self.dim = dim self.mode = mode # layers if mode == 'upsample2d': - self.resample = nn.Sequential( - Upsample(scale_factor=(2., 2.), mode='nearest-exact'), - nn.Conv2d(dim, dim // 2, 3, padding=1)) + self.resample = nn.Sequential(Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) elif mode == 'upsample3d': - self.resample = nn.Sequential( - Upsample(scale_factor=(2., 2.), mode='nearest-exact'), - nn.Conv2d(dim, dim // 2, 3, padding=1)) - self.time_conv = CausalConv3d( - dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + self.resample = nn.Sequential(Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) elif mode == 'downsample2d': - self.resample = nn.Sequential( - nn.ZeroPad2d((0, 1, 0, 1)), - nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) elif mode == 'downsample3d': - self.resample = nn.Sequential( - nn.ZeroPad2d((0, 1, 0, 1)), - nn.Conv2d(dim, dim, 3, stride=(2, 2))) - self.time_conv = CausalConv3d( - dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) else: self.resample = nn.Identity() @@ -110,21 +98,12 @@ class Resample(nn.Module): else: cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[ - idx] is not None and feat_cache[idx] != 'Rep': + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != 'Rep': # cache last frame of last two chunk - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) - if cache_x.shape[2] < 2 and feat_cache[ - idx] is not None and feat_cache[idx] == 'Rep': - cache_x = torch.cat([ - torch.zeros_like(cache_x).to(cache_x.device), - cache_x - ], + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == 'Rep': + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) if feat_cache[idx] == 'Rep': x = self.time_conv(x) else: @@ -133,8 +112,7 @@ class Resample(nn.Module): feat_idx[0] += 1 x = x.reshape(b, 2, c, t, h, w) - x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), - 3) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) x = x.reshape(b, c, t * 2, h, w) t = x.shape[2] x = rearrange(x, 'b c t h w -> (b t) c h w') @@ -154,8 +132,7 @@ class Resample(nn.Module): # # cache last frame of last two chunk # cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) - x = self.time_conv( - torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) feat_cache[idx] = cache_x feat_idx[0] += 1 return x @@ -192,11 +169,9 @@ class ResidualBlock(nn.Module): self.out_dim = out_dim # layers - self.residual = nn.Sequential( - RMS_norm(in_dim, images=False), nn.SiLU(), - CausalConv3d(in_dim, out_dim, 3, padding=1), - RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), - CausalConv3d(out_dim, out_dim, 3, padding=1)) + self.residual = nn.Sequential(RMS_norm(in_dim, images=False), nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), RMS_norm(out_dim, images=False), + nn.SiLU(), nn.Dropout(dropout), CausalConv3d(out_dim, out_dim, 3, padding=1)) self.shortcut = CausalConv3d(in_dim, out_dim, 1) \ if in_dim != out_dim else nn.Identity() @@ -208,10 +183,7 @@ class ResidualBlock(nn.Module): cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x @@ -244,10 +216,7 @@ class AttentionBlock(nn.Module): x = rearrange(x, 'b c t h w -> (b t) c h w') x = self.norm(x) # compute query, key, value - q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, - -1).permute(0, 1, 3, - 2).contiguous().chunk( - 3, dim=-1) + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1) # apply attention x = F.scaled_dot_product_attention( @@ -300,21 +269,18 @@ class Encoder3d(nn.Module): # downsample block if i != len(dim_mult) - 1: - mode = 'downsample3d' if temperal_downsample[ - i] else 'downsample2d' + mode = 'downsample3d' if temperal_downsample[i] else 'downsample2d' downsamples.append(Resample(out_dim, mode=mode)) scale /= 2.0 self.downsamples = nn.Sequential(*downsamples) # middle blocks - self.middle = nn.Sequential( - ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim), - ResidualBlock(out_dim, out_dim, dropout)) + self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout)) # output blocks - self.head = nn.Sequential( - RMS_norm(out_dim, images=False), nn.SiLU(), - CausalConv3d(out_dim, z_dim, 3, padding=1)) + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), CausalConv3d(out_dim, z_dim, 3, + padding=1)) def forward(self, x, feat_cache=None, feat_idx=[0]): if feat_cache is not None: @@ -322,11 +288,7 @@ class Encoder3d(nn.Module): cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv1(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 @@ -354,10 +316,7 @@ class Encoder3d(nn.Module): cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x @@ -393,9 +352,8 @@ class Decoder3d(nn.Module): self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) # middle blocks - self.middle = nn.Sequential( - ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]), - ResidualBlock(dims[0], dims[0], dropout)) + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) # upsample blocks upsamples = [] @@ -417,9 +375,7 @@ class Decoder3d(nn.Module): self.upsamples = nn.Sequential(*upsamples) # output blocks - self.head = nn.Sequential( - RMS_norm(out_dim, images=False), nn.SiLU(), - CausalConv3d(out_dim, 3, 3, padding=1)) + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), CausalConv3d(out_dim, 3, 3, padding=1)) def forward(self, x, feat_cache=None, feat_idx=[0]): ## conv1 @@ -428,11 +384,7 @@ class Decoder3d(nn.Module): cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv1(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 @@ -460,10 +412,7 @@ class Decoder3d(nn.Module): cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x @@ -501,12 +450,11 @@ class WanVAE_(nn.Module): self.temperal_upsample = temperal_downsample[::-1] # modules - self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, - attn_scales, self.temperal_downsample, dropout) + self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, + dropout) self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) self.conv2 = CausalConv3d(z_dim, z_dim, 1) - self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, - attn_scales, self.temperal_upsample, dropout) + self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout) def forward(self, x): mu, log_var = self.encode(x) @@ -522,20 +470,15 @@ class WanVAE_(nn.Module): for i in range(iter_): self._enc_conv_idx = [0] if i == 0: - out = self.encoder( - x[:, :, :1, :, :], - feat_cache=self._enc_feat_map, - feat_idx=self._enc_conv_idx) + out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) else: - out_ = self.encoder( - x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], - feat_cache=self._enc_feat_map, - feat_idx=self._enc_conv_idx) + out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) out = torch.cat([out, out_], 2) mu, log_var = self.conv1(out).chunk(2, dim=1) if isinstance(scale[0], torch.Tensor): - mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( - 1, self.z_dim, 1, 1, 1) + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1) else: mu = (mu - scale[0]) * scale[1] self.clear_cache() @@ -545,8 +488,7 @@ class WanVAE_(nn.Module): self.clear_cache() # z: [b,c,t,h,w] if isinstance(scale[0], torch.Tensor): - z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( - 1, self.z_dim, 1, 1, 1) + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1) else: z = z / scale[1] + scale[0] iter_ = z.shape[2] @@ -554,15 +496,9 @@ class WanVAE_(nn.Module): for i in range(iter_): self._conv_idx = [0] if i == 0: - out = self.decoder( - x[:, :, i:i + 1, :, :], - feat_cache=self._feat_map, - feat_idx=self._conv_idx) + out = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) else: - out_ = self.decoder( - x[:, :, i:i + 1, :, :], - feat_cache=self._feat_map, - feat_idx=self._conv_idx) + out_ = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) out = torch.cat([out, out_], 2) self.clear_cache() return out @@ -594,14 +530,13 @@ def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs): Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL. """ # params - cfg = dict( - dim=96, - z_dim=z_dim, - dim_mult=[1, 2, 4, 4], - num_res_blocks=2, - attn_scales=[], - temperal_downsample=[False, True, True], - dropout=0.0) + cfg = dict(dim=96, + z_dim=z_dim, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0) cfg.update(**kwargs) # init model @@ -610,29 +545,24 @@ def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs): # load checkpoint logging.info(f'loading {pretrained_path}') - model.load_state_dict( - torch.load(pretrained_path, map_location=device), assign=True) + model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True) return model class Wan2_1_VAE: - def __init__(self, - z_dim=16, - vae_pth='cache/vae_step_411000.pth', - dtype=torch.float, - device="cuda"): + def __init__(self, z_dim=16, vae_pth='cache/vae_step_411000.pth', dtype=torch.float, device="cuda"): self.dtype = dtype self.device = device mean = [ - -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, - 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 + -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, + -0.1922, -0.9497, 0.2503, -0.2921 ] std = [ - 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, - 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 + 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, + 1.1253, 2.8251, 1.9160 ] self.mean = torch.tensor(mean, dtype=dtype, device=device) self.std = torch.tensor(std, dtype=dtype, device=device) @@ -649,18 +579,11 @@ class Wan2_1_VAE: videos: A list of videos each with shape [C, T, H, W]. """ with amp.autocast(dtype=self.dtype): - return [ - self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0) - for u in videos - ] + return [self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0) for u in videos] def decode(self, zs): with amp.autocast(dtype=self.dtype): - return [ - self.model.decode(u.unsqueeze(0), - self.scale).float().clamp_(-1, 1).squeeze(0) - for u in zs - ] + return [self.model.decode(u.unsqueeze(0), self.scale).float().clamp_(-1, 1).squeeze(0) for u in zs] class LingBotWorld2WanVAE(nn.Module): @@ -681,8 +604,11 @@ class LingBotWorld2WanVAE(nn.Module): self.scale = [mean, 1.0 / std] if checkpoint_path is None: - self.model = WanVAE_(dim=96, z_dim=z_dim, dim_mult=[1, 2, 4, 4], - num_res_blocks=2, attn_scales=[], + self.model = WanVAE_(dim=96, + z_dim=z_dim, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], temperal_downsample=[False, True, True], dropout=0.0) else: diff --git a/fastvideo/models/vaes/ltx2vae.py b/fastvideo/models/vaes/ltx2vae.py index 545a7a1c3..8789d352a 100644 --- a/fastvideo/models/vaes/ltx2vae.py +++ b/fastvideo/models/vaes/ltx2vae.py @@ -229,9 +229,7 @@ def default_split_operation(length: int) -> DimensionIntervals: DEFAULT_SPLIT_OPERATION: SplitOperation = default_split_operation -def default_mapping_operation( - _intervals: DimensionIntervals, -) -> Tuple[List[slice], List[torch.Tensor | None]]: +def default_mapping_operation(_intervals: DimensionIntervals, ) -> Tuple[List[slice], List[torch.Tensor | None]]: return [slice(0, None)], [None] @@ -239,6 +237,7 @@ DEFAULT_MAPPING_OPERATION: MappingOperation = default_mapping_operation def split_in_spatial(size: int, overlap: int) -> SplitOperation: + def split(dimension_size: int) -> DimensionIntervals: if dimension_size <= size: return DEFAULT_SPLIT_OPERATION(dimension_size) @@ -291,6 +290,7 @@ def to_mapping_operation( map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor]], scale: int, ) -> MappingOperation: + def map_op(intervals: DimensionIntervals) -> Tuple[List[slice], List[torch.Tensor | None]]: output_slices: List[slice] = [] masks_1d: List[torch.Tensor | None] = [] @@ -330,13 +330,11 @@ def create_tiles_from_intervals_and_mappers( tile_out_coords = list(itertools.product(*full_dim_output_slices)) tile_mask_1ds = list(itertools.product(*full_dim_masks_1d)) for in_coord, out_coord, mask_1d in zip(tile_in_coords, tile_out_coords, tile_mask_1ds, strict=True): - tiles.append( - Tile( - in_coords=in_coord, - out_coords=out_coord, - masks_1d=mask_1d, - ) - ) + tiles.append(Tile( + in_coords=in_coord, + out_coords=out_coord, + masks_1d=mask_1d, + )) return tiles @@ -346,15 +344,11 @@ def create_tiles( mappers: List[MappingOperation], ) -> List[Tile]: if len(splitters) != len(latent_shape): - raise ValueError( - f"Number of splitters must be equal to number of dimensions in latent shape, " - f"got {len(splitters)} and {len(latent_shape)}" - ) + raise ValueError(f"Number of splitters must be equal to number of dimensions in latent shape, " + f"got {len(splitters)} and {len(latent_shape)}") if len(mappers) != len(latent_shape): - raise ValueError( - f"Number of mappers must be equal to number of dimensions in latent shape, " - f"got {len(mappers)} and {len(latent_shape)}" - ) + raise ValueError(f"Number of mappers must be equal to number of dimensions in latent shape, " + f"got {len(mappers)} and {len(latent_shape)}") intervals = [splitter(length) for splitter, length in zip(splitters, latent_shape, strict=True)] latent_intervals = LatentIntervals(original_shape=latent_shape, dimension_intervals=tuple(intervals)) return create_tiles_from_intervals_and_mappers(latent_intervals, mappers) @@ -671,17 +665,11 @@ class ResnetBlock3D(nn.Module): if inject_noise: self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1))) - self.conv_shortcut = ( - make_linear_nd(dims=dims, in_channels=in_channels, out_channels=out_channels) - if in_channels != out_channels - else nn.Identity() - ) + self.conv_shortcut = (make_linear_nd(dims=dims, in_channels=in_channels, out_channels=out_channels) + if in_channels != out_channels else nn.Identity()) - self.norm3 = ( - nn.GroupNorm(num_groups=1, num_channels=in_channels, eps=eps, affine=True) - if in_channels != out_channels - else nn.Identity() - ) + self.norm3 = (nn.GroupNorm(num_groups=1, num_channels=in_channels, eps=eps, affine=True) + if in_channels != out_channels else nn.Identity()) self.timestep_conditioning = timestep_conditioning @@ -718,15 +706,14 @@ class ResnetBlock3D(nn.Module): if timestep is None: raise ValueError("'timestep' must be provided when 'timestep_conditioning' is True") ada_values = self.scale_shift_table[None, ..., None, None, None].to( - device=hidden_states.device, dtype=hidden_states.dtype - ) + timestep.reshape( - batch_size, - 4, - -1, - timestep.shape[-3], - timestep.shape[-2], - timestep.shape[-1], - ) + device=hidden_states.device, dtype=hidden_states.dtype) + timestep.reshape( + batch_size, + 4, + -1, + timestep.shape[-3], + timestep.shape[-2], + timestep.shape[-1], + ) shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1) hidden_states = hidden_states * (1 + scale1) + shift1 @@ -766,18 +753,18 @@ class UNetMidBlock3D(nn.Module): """A 3D UNet mid-block with multiple residual blocks.""" def __init__( - self, - dims: int, - in_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_groups: int = 32, - norm_layer: NormLayerType = NormLayerType.GROUP_NORM, - inject_noise: bool = False, - timestep_conditioning: bool = False, - spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, - attention_head_dim: int | None = None, # unused, for compatibility + self, + dims: int, + in_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_groups: int = 32, + norm_layer: NormLayerType = NormLayerType.GROUP_NORM, + inject_noise: bool = False, + timestep_conditioning: bool = False, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + attention_head_dim: int | None = None, # unused, for compatibility ): super().__init__() resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) @@ -785,27 +772,23 @@ class UNetMidBlock3D(nn.Module): self.timestep_conditioning = timestep_conditioning if timestep_conditioning: - self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings( - embedding_dim=in_channels * 4, size_emb_dim=0 - ) + self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=in_channels * 4, + size_emb_dim=0) - self.res_blocks = nn.ModuleList( - [ - ResnetBlock3D( - dims=dims, - in_channels=in_channels, - out_channels=in_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - norm_layer=norm_layer, - inject_noise=inject_noise, - timestep_conditioning=timestep_conditioning, - spatial_padding_mode=spatial_padding_mode, - ) - for _ in range(num_layers) - ] - ) + self.res_blocks = nn.ModuleList([ + ResnetBlock3D( + dims=dims, + in_channels=in_channels, + out_channels=in_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + norm_layer=norm_layer, + inject_noise=inject_noise, + timestep_conditioning=timestep_conditioning, + spatial_padding_mode=spatial_padding_mode, + ) for _ in range(num_layers) + ]) def forward( self, @@ -866,9 +849,7 @@ class Timesteps(nn.Module): def forward(self, timesteps: torch.Tensor) -> torch.Tensor: half_dim = self.num_channels // 2 - exponent = -math.log(10000) * torch.arange( - start=0, end=half_dim, dtype=torch.float32, device=timesteps.device - ) + exponent = -math.log(10000) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device) exponent = exponent / (half_dim - self.downscale_freq_shift) emb = torch.exp(exponent) emb = timesteps[:, None].float() * emb[None, :] @@ -1267,9 +1248,7 @@ class VideoEncoder(nn.Module): self.down_blocks.append(block) if norm_layer == NormLayerType.GROUP_NORM: - self.conv_norm_out = nn.GroupNorm( - num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6 - ) + self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6) elif norm_layer == NormLayerType.PIXEL_NORM: self.conv_norm_out = PixelNorm() @@ -1296,10 +1275,8 @@ class VideoEncoder(nn.Module): def forward(self, sample: torch.Tensor) -> torch.Tensor: frames_count = sample.shape[2] if ((frames_count - 1) % 8) != 0: - raise ValueError( - "Invalid number of frames: Encode input must have 1 + 8 * x frames " - "(e.g., 1, 9, 17, ...). Please check your input." - ) + raise ValueError("Invalid number of frames: Encode input must have 1 + 8 * x frames " + "(e.g., 1, 9, 17, ...). Please check your input.") sample = patchify(sample, patch_size_hw=self.patch_size, patch_size_t=1) sample = self.conv_in(sample) @@ -1313,9 +1290,7 @@ class VideoEncoder(nn.Module): if self.latent_log_var == LogVarianceType.UNIFORM: if sample.shape[1] < 2: - raise ValueError( - f"Invalid channel count for UNIFORM mode: expected at least 2, got {sample.shape[1]}" - ) + raise ValueError(f"Invalid channel count for UNIFORM mode: expected at least 2, got {sample.shape[1]}") means = sample[:, :-1, ...] logvar = sample[:, -1:, ...] num_channels = means.shape[1] @@ -1413,9 +1388,7 @@ class VideoDecoder(nn.Module): self.up_blocks.append(block) if norm_layer == NormLayerType.GROUP_NORM: - self.conv_norm_out = nn.GroupNorm( - num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6 - ) + self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6) elif norm_layer == NormLayerType.PIXEL_NORM: self.conv_norm_out = PixelNorm() @@ -1432,9 +1405,8 @@ class VideoDecoder(nn.Module): if timestep_conditioning: self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0)) - self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings( - embedding_dim=feature_channels * 2, size_emb_dim=0 - ) + self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=feature_channels * 2, + size_emb_dim=0) self.last_scale_shift_table = nn.Parameter(torch.empty(2, feature_channels)) def forward( @@ -1446,21 +1418,18 @@ class VideoDecoder(nn.Module): batch_size = sample.shape[0] if self.timestep_conditioning: - noise = ( - torch.randn( - sample.size(), - generator=generator, - dtype=sample.dtype, - device=sample.device, - ) - * self.decode_noise_scale - ) + noise = (torch.randn( + sample.size(), + generator=generator, + dtype=sample.dtype, + device=sample.device, + ) * self.decode_noise_scale) sample = noise + (1.0 - self.decode_noise_scale) * sample sample = self.per_channel_statistics.un_normalize(sample) if timestep is None and self.timestep_conditioning: - timestep = torch.full((batch_size,), self.decode_timestep, device=sample.device, dtype=sample.dtype) + timestep = torch.full((batch_size, ), self.decode_timestep, device=sample.device, dtype=sample.dtype) sample = self.conv_in(sample, causal=self.causal) @@ -1492,15 +1461,14 @@ class VideoDecoder(nn.Module): ) embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1) ada_values = self.last_scale_shift_table[None, ..., None, None, None].to( - device=sample.device, dtype=sample.dtype - ) + embedded_timestep.reshape( - batch_size, - 2, - -1, - embedded_timestep.shape[-3], - embedded_timestep.shape[-2], - embedded_timestep.shape[-1], - ) + device=sample.device, dtype=sample.dtype) + embedded_timestep.reshape( + batch_size, + 2, + -1, + embedded_timestep.shape[-3], + embedded_timestep.shape[-2], + embedded_timestep.shape[-1], + ) shift, scale = ada_values.unbind(dim=1) sample = sample * (1 + scale) + shift @@ -1631,7 +1599,8 @@ class LTX2CausalVideoAutoencoder(nn.Module): self.enable_channels_last_3d() def _as_channels_last_3d(self, tensor: torch.Tensor) -> torch.Tensor: - if not self._use_channels_last_3d or tensor.ndim != 5 or tensor.is_contiguous(memory_format=torch.channels_last_3d): + if not self._use_channels_last_3d or tensor.ndim != 5 or tensor.is_contiguous( + memory_format=torch.channels_last_3d): return tensor return tensor.contiguous(memory_format=torch.channels_last_3d) @@ -1783,9 +1752,7 @@ class LTX2CausalVideoAutoencoder(nn.Module): Yields: Video chunks (B, C, T, H, W) by temporal slices. """ - full_video_shape = VideoLatentShape.from_torch_shape(latent.shape).upscale( - self.TIME_SCALE, self.SPATIAL_SCALE - ) + full_video_shape = VideoLatentShape.from_torch_shape(latent.shape).upscale(self.TIME_SCALE, self.SPATIAL_SCALE) tiles = self._prepare_tiles(latent, tiling_config) temporal_groups = self._group_tiles_by_temporal_slice(tiles) @@ -1796,9 +1763,8 @@ class LTX2CausalVideoAutoencoder(nn.Module): for temporal_group_tiles in temporal_groups: curr_temporal_slice = temporal_group_tiles[0].out_coords[2] - temporal_tile_buffer_shape = full_video_shape._replace( - frames=curr_temporal_slice.stop - curr_temporal_slice.start, - ) + temporal_tile_buffer_shape = full_video_shape._replace(frames=curr_temporal_slice.stop - + curr_temporal_slice.start, ) buffer = torch.zeros( temporal_tile_buffer_shape.to_torch_shape(), @@ -1821,14 +1787,12 @@ class LTX2CausalVideoAutoencoder(nn.Module): temporal_overlap_slice = slice(curr_temporal_slice.start - previous_temporal_slice.start, None) previous_chunk[:, :, temporal_overlap_slice, :, :] += buffer[:, :, slice(0, overlap_len), :, :] - previous_weights[:, :, temporal_overlap_slice, :, :] += curr_weights[ - :, :, slice(0, overlap_len), :, : - ] + previous_weights[:, :, temporal_overlap_slice, :, :] += curr_weights[:, :, + slice(0, overlap_len), :, :] buffer[:, :, slice(0, overlap_len), :, :] = previous_chunk[:, :, temporal_overlap_slice, :, :] - curr_weights[:, :, slice(0, overlap_len), :, :] = previous_weights[ - :, :, temporal_overlap_slice, :, : - ] + curr_weights[:, :, slice(0, overlap_len), :, :] = previous_weights[:, :, + temporal_overlap_slice, :, :] # Yield the non-overlapping part of the previous chunk previous_weights = previous_weights.clamp(min=1e-8) @@ -1845,5 +1809,6 @@ class LTX2CausalVideoAutoencoder(nn.Module): previous_weights = previous_weights.clamp(min=1e-8) yield previous_chunk / previous_weights + # Entry point for model registry EntryClass = LTX2CausalVideoAutoencoder diff --git a/fastvideo/models/vaes/minimax_h3_audio.py b/fastvideo/models/vaes/minimax_h3_audio.py index 8b07dd731..6aefc411c 100644 --- a/fastvideo/models/vaes/minimax_h3_audio.py +++ b/fastvideo/models/vaes/minimax_h3_audio.py @@ -1,6 +1,5 @@ # Copyright 2025 MiniMax authors and HuggingFace Team # SPDX-License-Identifier: Apache-2.0 - """Native MiniMax H3 waveform autoencoder.""" from __future__ import annotations @@ -74,7 +73,7 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> if attenuation > 50.0: beta = 0.1102 * (attenuation - 8.7) elif attenuation >= 21.0: - beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0) + beta = 0.5842 * (attenuation - 21)**0.4 + 0.07886 * (attenuation - 21.0) else: beta = 0.0 @@ -284,9 +283,8 @@ class MiniMaxH3AudioCausalAttention(nn.Module): self.qkv.weight, torch.cat((self.q_bias, self.zero_k_bias, self.v_bias)), ) - query, key, value = ( - qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim).permute(2, 0, 1, 3, 4).unbind(0) - ) + query, key, value = (qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim).permute(2, 0, 1, 3, + 4).unbind(0)) # H3's projection is an unusual causal flat stream, not a VAE spatial-attention block. hidden_states = F.scaled_dot_product_attention( @@ -325,13 +323,10 @@ class MiniMaxH3AudioAMPBlock(nn.Module): _wn_conv1d(channels, channels, kernel_size, dilation=d, padding=(kernel_size * d - d) // 2) for d in dilation ]) - self.convs2 = nn.ModuleList([ - _wn_conv1d(channels, channels, kernel_size, dilation=1, padding=(kernel_size - 1) // 2) - for _ in dilation - ]) + self.convs2 = nn.ModuleList( + [_wn_conv1d(channels, channels, kernel_size, dilation=1, padding=(kernel_size - 1) // 2) for _ in dilation]) self.activations = nn.ModuleList([ - MiniMaxH3AudioActivation1d(activation=MiniMaxH3AudioSnakeBeta(channels)) - for _ in range(2 * len(dilation)) + MiniMaxH3AudioActivation1d(activation=MiniMaxH3AudioSnakeBeta(channels)) for _ in range(2 * len(dilation)) ]) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -361,19 +356,21 @@ class MiniMaxH3AudioBigVGANDecoder(nn.Module): self.ups = nn.ModuleList() for i, (rate, kernel) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): - self.ups.append(nn.ModuleList([ - weight_norm(nn.ConvTranspose1d( - upsample_initial_channel // (2**i), - upsample_initial_channel // (2 ** (i + 1)), - kernel, - rate, - padding=(kernel - rate) // 2, - )) - ])) + self.ups.append( + nn.ModuleList([ + weight_norm( + nn.ConvTranspose1d( + upsample_initial_channel // (2**i), + upsample_initial_channel // (2**(i + 1)), + kernel, + rate, + padding=(kernel - rate) // 2, + )) + ])) self.resblocks = nn.ModuleList() for i in range(self.num_upsamples): - channels = upsample_initial_channel // (2 ** (i + 1)) + channels = upsample_initial_channel // (2**(i + 1)) for kernel, dilation in zip(resblock_kernel_sizes, resblock_dilation_sizes): self.resblocks.append(MiniMaxH3AudioAMPBlock(channels, kernel, tuple(dilation))) @@ -424,15 +421,11 @@ class MiniMaxH3AudioVAE(nn.Module): ) if math.prod(decoder_rates) != self.hop_length: - raise ValueError( - f"`decoder_rates` must upsample by the encoder hop length {self.hop_length}, got " - f"{math.prod(decoder_rates)}." - ) + raise ValueError(f"`decoder_rates` must upsample by the encoder hop length {self.hop_length}, got " + f"{math.prod(decoder_rates)}.") if arch.latent_dim % arch.latent_channels != 0: - raise ValueError( - f"`latent_dim` ({arch.latent_dim}) must be a multiple of `latent_channels` " - f"({arch.latent_channels})." - ) + raise ValueError(f"`latent_dim` ({arch.latent_dim}) must be a multiple of `latent_channels` " + f"({arch.latent_channels}).") self.encoder = MiniMaxH3AudioEncoder( d_model=arch.encoder_dim, @@ -455,8 +448,7 @@ class MiniMaxH3AudioVAE(nn.Module): upsample_kernel_sizes=tuple(int(kernel) for kernel in arch.decoder_kernel_sizes), resblock_kernel_sizes=tuple(int(kernel) for kernel in arch.resblock_kernel_sizes), resblock_dilation_sizes=tuple( - tuple(int(dilation) for dilation in group) for group in arch.resblock_dilation_sizes - ), + tuple(int(dilation) for dilation in group) for group in arch.resblock_dilation_sizes), ) # The H3 checkpoint and waveform numerics require this component to stay FP32. @@ -485,7 +477,7 @@ class MiniMaxH3AudioVAE(nn.Module): posterior = MiniMaxH3AudioDiagonalGaussianDistribution(mean, logs) if not return_dict: - return (posterior,) + return (posterior, ) return MiniMaxH3AudioEncoderOutput(latent_dist=posterior) def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor: @@ -501,8 +493,7 @@ class MiniMaxH3AudioVAE(nn.Module): ) -> MiniMaxH3AudioDecoderOutput | tuple[torch.Tensor]: if latents.ndim != 3: raise ValueError( - f"`latents` must have shape [batch_size, latent_channels, num_frames], got {tuple(latents.shape)}." - ) + f"`latents` must have shape [batch_size, latent_channels, num_frames], got {tuple(latents.shape)}.") decoder_dtype = _module_dtype(self.decoder) decoded = self.decoder(self.dec_in_proj(latents.to(decoder_dtype))) @@ -510,7 +501,7 @@ class MiniMaxH3AudioVAE(nn.Module): decoded = decoded.float() if not return_dict: - return (decoded,) + return (decoded, ) return MiniMaxH3AudioDecoderOutput(sample=decoded) def forward( diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py index d925a70e5..d4c6b95a6 100644 --- a/fastvideo/models/vaes/minimax_h3_video.py +++ b/fastvideo/models/vaes/minimax_h3_video.py @@ -101,6 +101,7 @@ class MiniMaxH3VideoGroupNorm(nn.GroupNorm): class MiniMaxH3VideoResnetBlock3d(nn.Module): + def __init__( self, in_channels: int, @@ -144,6 +145,7 @@ class MiniMaxH3VideoResnetBlock3d(nn.Module): class MiniMaxH3VideoDownsample3d(nn.Module): + def __init__( self, in_channels: int, @@ -171,6 +173,7 @@ class MiniMaxH3VideoDownsample3d(nn.Module): class MiniMaxH3VideoDownBlock3d(nn.Module): + def __init__( self, in_channels: int, @@ -183,31 +186,26 @@ class MiniMaxH3VideoDownBlock3d(nn.Module): spatial_padding_mode: str = "reflect", ) -> None: super().__init__() - self.resnets = nn.ModuleList( - [ - MiniMaxH3VideoResnetBlock3d( - in_channels=in_channels if index == 0 else out_channels, - out_channels=out_channels, - norm_num_groups=norm_num_groups, - norm_eps=norm_eps, - spatial_padding_mode=spatial_padding_mode, - ) - for index in range(num_layers) - ] - ) + self.resnets = nn.ModuleList([ + MiniMaxH3VideoResnetBlock3d( + in_channels=in_channels if index == 0 else out_channels, + out_channels=out_channels, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + spatial_padding_mode=spatial_padding_mode, + ) for index in range(num_layers) + ]) self.downsamplers = None if temporal_downsample_factor * spatial_downsample_factor > 1: - self.downsamplers = nn.ModuleList( - [ - MiniMaxH3VideoDownsample3d( - out_channels, - out_channels, - temporal_stride=temporal_downsample_factor, - spatial_stride=spatial_downsample_factor, - spatial_padding_mode=spatial_padding_mode, - ) - ] - ) + self.downsamplers = nn.ModuleList([ + MiniMaxH3VideoDownsample3d( + out_channels, + out_channels, + temporal_stride=temporal_downsample_factor, + spatial_stride=spatial_downsample_factor, + spatial_padding_mode=spatial_padding_mode, + ) + ]) self.gradient_checkpointing = False def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -223,6 +221,7 @@ class MiniMaxH3VideoDownBlock3d(nn.Module): class MiniMaxH3VideoEncoder3d(nn.Module): + def __init__( self, in_channels: int, @@ -244,22 +243,19 @@ class MiniMaxH3VideoEncoder3d(nn.Module): temporal_padding=2, spatial_padding_mode=spatial_padding_mode, ) - block_in_channels = (block_out_channels[0],) + tuple(block_out_channels[:-1]) - self.down_blocks = nn.ModuleList( - [ - MiniMaxH3VideoDownBlock3d( - in_channels=block_in_channels[index], - out_channels=block_out_channels[index], - num_layers=layers_per_block, - temporal_downsample_factor=temporal_downsample_factors[index], - spatial_downsample_factor=spatial_downsample_factors[index], - norm_num_groups=norm_num_groups, - norm_eps=norm_eps, - spatial_padding_mode=spatial_padding_mode, - ) - for index in range(len(block_out_channels)) - ] - ) + block_in_channels = (block_out_channels[0], ) + tuple(block_out_channels[:-1]) + self.down_blocks = nn.ModuleList([ + MiniMaxH3VideoDownBlock3d( + in_channels=block_in_channels[index], + out_channels=block_out_channels[index], + num_layers=layers_per_block, + temporal_downsample_factor=temporal_downsample_factors[index], + spatial_downsample_factor=spatial_downsample_factors[index], + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + spatial_padding_mode=spatial_padding_mode, + ) for index in range(len(block_out_channels)) + ]) self.norm_out = MiniMaxH3VideoGroupNorm(norm_num_groups, block_out_channels[-1], eps=norm_eps, affine=True) self.conv_out = MiniMaxH3VideoCausalConv3d( block_out_channels[-1], @@ -278,11 +274,12 @@ class MiniMaxH3VideoEncoder3d(nn.Module): class MiniMaxH3VideoRotaryPosEmbed(nn.Module): + def __init__(self, dim: int, theta: float = 100.0, num_axes: int = 3) -> None: super().__init__() if dim % (2 * num_axes) != 0: raise ValueError(f"`dim` {dim} must be divisible by `2 * num_axes` {2 * num_axes}.") - inv_freq = 1.0 / theta ** torch.arange(0, 1, 2 * num_axes / dim, dtype=torch.float32) + inv_freq = 1.0 / theta**torch.arange(0, 1, 2 * num_axes / dim, dtype=torch.float32) self.register_buffer("inv_freq", inv_freq, persistent=False) def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -292,6 +289,7 @@ class MiniMaxH3VideoRotaryPosEmbed(nn.Module): class MiniMaxH3VideoAttention(nn.Module): + def __init__(self, dim: int, heads: int, dim_head: int, eps: float = 1e-5, bias: bool = True) -> None: super().__init__() self.heads = heads @@ -337,6 +335,7 @@ class MiniMaxH3VideoAttention(nn.Module): class MiniMaxH3VideoSwiGLU(nn.Module): + def __init__(self, dim_in: int, dim_out: int, bias: bool = True) -> None: super().__init__() self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias) @@ -348,16 +347,15 @@ class MiniMaxH3VideoSwiGLU(nn.Module): class MiniMaxH3VideoFeedForward(nn.Module): + def __init__(self, dim: int, mult: int = 4, bias: bool = True) -> None: super().__init__() inner_dim = int(dim * mult) - self.net = nn.ModuleList( - [ - MiniMaxH3VideoSwiGLU(dim, inner_dim, bias=bias), - nn.Dropout(0.0), - nn.Linear(inner_dim, dim, bias=bias), - ] - ) + self.net = nn.ModuleList([ + MiniMaxH3VideoSwiGLU(dim, inner_dim, bias=bias), + nn.Dropout(0.0), + nn.Linear(inner_dim, dim, bias=bias), + ]) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: for module in self.net: @@ -366,6 +364,7 @@ class MiniMaxH3VideoFeedForward(nn.Module): class MiniMaxH3VideoTransformerBlock(nn.Module): + def __init__( self, dim: int, @@ -395,6 +394,7 @@ class MiniMaxH3VideoTransformerBlock(nn.Module): class MiniMaxH3VideoViTDecoder3d(nn.Module): + def __init__( self, in_channels: int, @@ -419,18 +419,15 @@ class MiniMaxH3VideoViTDecoder3d(nn.Module): self.rope = MiniMaxH3VideoRotaryPosEmbed(int(attention_head_dim * rope_dim_ratio), theta=rope_theta) self.proj_in = nn.Linear(in_channels, dim) self.register_tokens = nn.Parameter(torch.zeros(1, num_register_tokens, dim)) - self.transformer_blocks = nn.ModuleList( - [ - MiniMaxH3VideoTransformerBlock( - dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - ffn_mult=ffn_mult, - eps=norm_eps, - ) - for _ in range(num_layers) - ] - ) + self.transformer_blocks = nn.ModuleList([ + MiniMaxH3VideoTransformerBlock( + dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + ffn_mult=ffn_mult, + eps=norm_eps, + ) for _ in range(num_layers) + ]) self.norm_out = nn.LayerNorm(dim, elementwise_affine=True, eps=norm_eps) self.proj_out = nn.Linear(dim, out_channels * patch_size_t * patch_size * patch_size) self.gradient_checkpointing = False @@ -672,9 +669,9 @@ class AutoencoderKLMiniMaxH3(nn.Module): if column_index > 0: tile = self._blend(row[column_index - 1], tile, width_overlaps[column_index - 1], -1) if row_index < len(tiles) - 1: - tile = tile[..., : -height_overlaps[row_index], :] + tile = tile[..., :-height_overlaps[row_index], :] if column_index < len(row) - 1: - tile = tile[..., :, : -width_overlaps[column_index]] + tile = tile[..., :, :-width_overlaps[column_index]] result_row.append(tile) result_rows.append(torch.cat(result_row, dim=-1)) return torch.cat(result_rows, dim=-2) @@ -783,12 +780,8 @@ class AutoencoderKLMiniMaxH3(nn.Module): if pad_tokens > 0: intra_tail = self.config.clip_length % temporal_ratio num_tokens_before_pad = z.shape[2] - pad_tokens - pad_frames = sum( - intra_tail - if intra_tail and (num_tokens_before_pad + offset) % tokens_chunk_size == 0 - else temporal_ratio - for offset in range(pad_tokens) - ) + pad_frames = sum(intra_tail if intra_tail and (num_tokens_before_pad + offset) % + tokens_chunk_size == 0 else temporal_ratio for offset in range(pad_tokens)) decoded = decoded[:, :, :-pad_frames] return decoded @@ -803,7 +796,7 @@ class AutoencoderKLMiniMaxH3(nn.Module): moments = self._encode(x) posterior = DiagonalGaussianDistribution(moments) if not return_dict: - return (posterior,) + return (posterior, ) return AutoencoderKLOutput(latent_dist=posterior) def encode_keyframe( @@ -820,7 +813,7 @@ class AutoencoderKLMiniMaxH3(nn.Module): moments = self._encode_clip(x) posterior = DiagonalGaussianDistribution(moments) if not return_dict: - return (posterior,) + return (posterior, ) return AutoencoderKLOutput(latent_dist=posterior) def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]: @@ -829,7 +822,7 @@ class AutoencoderKLMiniMaxH3(nn.Module): else: decoded = self._decode(z) if not return_dict: - return (decoded,) + return (decoded, ) return DecoderOutput(sample=decoded) def forward( diff --git a/fastvideo/models/vaes/oobleck.py b/fastvideo/models/vaes/oobleck.py index 8ba2fd7c8..2c55b25e6 100644 --- a/fastvideo/models/vaes/oobleck.py +++ b/fastvideo/models/vaes/oobleck.py @@ -47,12 +47,17 @@ class Snake1d(nn.Module): class OobleckResidualUnit(nn.Module): + def __init__(self, dimension: int = 16, dilation: int = 1): super().__init__() pad = ((7 - 1) * dilation) // 2 self.snake1 = Snake1d(dimension) self.conv1 = weight_norm(nn.Conv1d( - dimension, dimension, kernel_size=7, dilation=dilation, padding=pad, + dimension, + dimension, + kernel_size=7, + dilation=dilation, + padding=pad, )) self.snake2 = Snake1d(dimension) self.conv2 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=1)) @@ -67,17 +72,21 @@ class OobleckResidualUnit(nn.Module): class OobleckEncoderBlock(nn.Module): + def __init__(self, input_dim: int, output_dim: int, stride: int = 1): super().__init__() self.res_unit1 = OobleckResidualUnit(input_dim, dilation=1) self.res_unit2 = OobleckResidualUnit(input_dim, dilation=3) self.res_unit3 = OobleckResidualUnit(input_dim, dilation=9) self.snake1 = Snake1d(input_dim) - self.conv1 = weight_norm(nn.Conv1d( - input_dim, output_dim, - kernel_size=2 * stride, stride=stride, - padding=math.ceil(stride / 2), - )) + self.conv1 = weight_norm( + nn.Conv1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + )) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.res_unit1(x) @@ -87,14 +96,18 @@ class OobleckEncoderBlock(nn.Module): class OobleckDecoderBlock(nn.Module): + def __init__(self, input_dim: int, output_dim: int, stride: int = 1): super().__init__() self.snake1 = Snake1d(input_dim) - self.conv_t1 = weight_norm(nn.ConvTranspose1d( - input_dim, output_dim, - kernel_size=2 * stride, stride=stride, - padding=math.ceil(stride / 2), - )) + self.conv_t1 = weight_norm( + nn.ConvTranspose1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + )) self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1) self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3) self.res_unit3 = OobleckResidualUnit(output_dim, dilation=9) @@ -120,8 +133,10 @@ class OobleckDiagonalGaussianDistribution: def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: noise = torch.randn( - self.mean.shape, generator=generator, - device=self.parameters.device, dtype=self.parameters.dtype, + self.mean.shape, + generator=generator, + device=self.parameters.device, + dtype=self.parameters.dtype, ) return self.mean + self.std * noise @@ -135,6 +150,7 @@ class OobleckDecoderOutput: class OobleckEncoder(nn.Module): + def __init__( self, encoder_hidden_size: int, @@ -146,20 +162,25 @@ class OobleckEncoder(nn.Module): strides = downsampling_ratios channel_multiples = [1] + list(channel_multiples) self.conv1 = weight_norm(nn.Conv1d( - audio_channels, encoder_hidden_size, kernel_size=7, padding=3, + audio_channels, + encoder_hidden_size, + kernel_size=7, + padding=3, )) self.block = nn.ModuleList([ OobleckEncoderBlock( input_dim=encoder_hidden_size * channel_multiples[i], output_dim=encoder_hidden_size * channel_multiples[i + 1], stride=s, - ) - for i, s in enumerate(strides) + ) for i, s in enumerate(strides) ]) d_model = encoder_hidden_size * channel_multiples[-1] self.snake1 = Snake1d(d_model) self.conv2 = weight_norm(nn.Conv1d( - d_model, encoder_hidden_size, kernel_size=3, padding=1, + d_model, + encoder_hidden_size, + kernel_size=3, + padding=1, )) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -171,6 +192,7 @@ class OobleckEncoder(nn.Module): class OobleckDecoder(nn.Module): + def __init__( self, channels: int, @@ -183,20 +205,25 @@ class OobleckDecoder(nn.Module): strides = upsampling_ratios channel_multiples = [1] + list(channel_multiples) self.conv1 = weight_norm(nn.Conv1d( - input_channels, channels * channel_multiples[-1], - kernel_size=7, padding=3, + input_channels, + channels * channel_multiples[-1], + kernel_size=7, + padding=3, )) self.block = nn.ModuleList([ OobleckDecoderBlock( input_dim=channels * channel_multiples[len(strides) - i], output_dim=channels * channel_multiples[len(strides) - i - 1], stride=s, - ) - for i, s in enumerate(strides) + ) for i, s in enumerate(strides) ]) self.snake1 = Snake1d(channels) self.conv2 = weight_norm(nn.Conv1d( - channels, audio_channels, kernel_size=7, padding=3, bias=False, + channels, + audio_channels, + kernel_size=7, + padding=3, + bias=False, )) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -270,7 +297,8 @@ class OobleckVAE(nn.Module): ) def encode( - self, x: torch.Tensor, + self, + x: torch.Tensor, ) -> OobleckDiagonalGaussianDistribution: return OobleckDiagonalGaussianDistribution(self.encoder(x)) @@ -278,7 +306,9 @@ class OobleckVAE(nn.Module): return OobleckDecoderOutput(sample=self.decoder(z)) def forward( - self, sample: torch.Tensor, sample_posterior: bool = False, + self, + sample: torch.Tensor, + sample_posterior: bool = False, ) -> OobleckDecoderOutput: posterior = self.encode(sample) z = posterior.sample() if sample_posterior else posterior.mode() @@ -317,7 +347,9 @@ class OobleckVAE(nn.Module): from huggingface_hub import snapshot_download allow = ["vae/*"] if subfolder else ["*"] root = snapshot_download( - repo_id=model_path, token=resolve_hf_token(), allow_patterns=allow, + repo_id=model_path, + token=resolve_hf_token(), + allow_patterns=allow, ) if subfolder: root = os.path.join(root, subfolder) @@ -326,10 +358,8 @@ class OobleckVAE(nn.Module): cfg_path = os.path.join(root, "config.json") if not os.path.isfile(cfg_path): - raise FileNotFoundError( - f"Expected config.json at {cfg_path}. If using a HF repo, " - f"pass subfolder='vae'." - ) + raise FileNotFoundError(f"Expected config.json at {cfg_path}. If using a HF repo, " + f"pass subfolder='vae'.") with open(cfg_path) as f: cfg = json.load(f) cfg_fields = {k: v for k, v in cfg.items() if not k.startswith("_")} @@ -347,24 +377,23 @@ class OobleckVAE(nn.Module): if os.path.isfile(alt): weights_path = alt else: - raise FileNotFoundError( - f"No safetensors weights under {root}. Expected " - f"diffusion_pytorch_model.safetensors." - ) + raise FileNotFoundError(f"No safetensors weights under {root}. Expected " + f"diffusion_pytorch_model.safetensors.") state = load_file(weights_path) missing, unexpected = model.load_state_dict(state, strict=False) if missing: - raise RuntimeError( - f"OobleckVAE missing {len(missing)} keys from {weights_path}: " - f"{missing[:5]}" - ) + raise RuntimeError(f"OobleckVAE missing {len(missing)} keys from {weights_path}: " + f"{missing[:5]}") if unexpected: # Non-critical: some checkpoints embed the VAE inside a larger # container (e.g. `pretransform.model.*`). Log the count so # genuine loader regressions don't go unnoticed. logger.debug( "OobleckVAE: ignored %d unexpected keys from %s " - "(first 3: %s)", len(unexpected), weights_path, unexpected[:3], + "(first 3: %s)", + len(unexpected), + weights_path, + unexpected[:3], ) if torch_dtype is not None: diff --git a/fastvideo/models/vaes/sa_audio.py b/fastvideo/models/vaes/sa_audio.py index 1031f9533..c75e3b86b 100644 --- a/fastvideo/models/vaes/sa_audio.py +++ b/fastvideo/models/vaes/sa_audio.py @@ -50,10 +50,8 @@ class SAAudioVAEModel(nn.Module): path = self.pretrained_path if not path: - raise ValueError( - "OobleckVAEConfig.pretrained_path must be set; expected " - "`stabilityai/stable-audio-open-1.0` or a local path." - ) + raise ValueError("OobleckVAEConfig.pretrained_path must be set; expected " + "`stabilityai/stable-audio-open-1.0` or a local path.") dtype = getattr(torch, self.pretrained_dtype, torch.float32) # If the caller already pointed us at the VAE dir directly, drop # the subfolder. Otherwise pass through (default "vae"). diff --git a/fastvideo/models/vaes/wanvae.py b/fastvideo/models/vaes/wanvae.py index b3eef55a4..63403130b 100644 --- a/fastvideo/models/vaes/wanvae.py +++ b/fastvideo/models/vaes/wanvae.py @@ -24,8 +24,7 @@ from einops import rearrange from fastvideo.configs.models.vaes import WanVAEConfig from fastvideo.layers.activation import get_act_fn -from fastvideo.models.vaes.common import (DiagonalGaussianDistribution, - ParallelTiledVAE) +from fastvideo.models.vaes.common import (DiagonalGaussianDistribution, ParallelTiledVAE) CACHE_T = 2 @@ -56,7 +55,9 @@ def forward_context(first_frame_arg=False, first_chunk.reset(first_chunk_token) use_light_vae.reset(use_light_vae_token) + class AvgDown3D(nn.Module): + def __init__( self, in_channels, @@ -110,6 +111,7 @@ class AvgDown3D(nn.Module): class DupUp3D(nn.Module): + def __init__( self, in_channels: int, @@ -151,9 +153,10 @@ class DupUp3D(nn.Module): _first_chunk = first_chunk.get() if _first_chunk: - x = x[:, :, self.factor_t - 1 :, :, :] + x = x[:, :, self.factor_t - 1:, :, :] return x + class WanCausalConv3d(nn.Conv3d): r""" A custom 3D causal convolution layer with feature caching support. @@ -186,8 +189,7 @@ class WanCausalConv3d(nn.Conv3d): ) self.padding: tuple[int, int, int] # Set up causal padding - self._padding: tuple[int, ...] = (self.padding[2], self.padding[2], - self.padding[1], self.padding[1], + self._padding: tuple[int, ...] = (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0) self.padding = (0, 0, 0) @@ -200,8 +202,8 @@ class WanCausalConv3d(nn.Conv3d): x = torch.cat([cache_x, x], dim=2) padding[4] -= cache_x.shape[2] x = F.pad(x, padding) - x = x.to(self.weight.dtype) if current_platform.is_mps( - ) else x # casting needed for mps since amp isn't supported + x = x.to( + self.weight.dtype) if current_platform.is_mps() else x # casting needed for mps since amp isn't supported return super().forward(x) @@ -217,11 +219,7 @@ class WanRMS_norm(nn.Module): bias (bool, optional): Whether to include a learnable bias term. Default is False. """ - def __init__(self, - dim: int, - channel_first: bool = True, - images: bool = True, - bias: bool = False) -> None: + def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None: super().__init__() broadcastable_dims = (1, 1, 1) if not images else (1, 1) shape = (dim, *broadcastable_dims) if channel_first else (dim, ) @@ -232,8 +230,7 @@ class WanRMS_norm(nn.Module): self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 def forward(self, x): - return F.normalize(x, dim=(1 if self.channel_first else - -1)) * self.scale * self.gamma + self.bias + return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias class WanUpsample(nn.Upsample): @@ -285,20 +282,13 @@ class WanResample(nn.Module): WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, upsample_out_dim, 3, padding=1), ) - self.time_conv = WanCausalConv3d(dim, - dim * 2, (3, 1, 1), - padding=(1, 0, 0)) + self.time_conv = WanCausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) elif mode == "downsample2d": - self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), - nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) elif mode == "downsample3d": - self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), - nn.Conv2d(dim, dim, 3, stride=(2, 2))) - self.time_conv = WanCausalConv3d(dim, - dim, (3, 1, 1), - stride=(2, 1, 1), - padding=(0, 0, 0)) + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = WanCausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) else: self.resample = nn.Identity() @@ -325,11 +315,10 @@ class WanResample(nn.Module): else: if cache_x.shape[2] < 2: cache_x = torch.cat([ - _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - x.device - ), + _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(x.device), cache_x, - ], dim=2) + ], + dim=2) x = self.time_conv(x, _feat_cache[idx]) _feat_cache[idx] = cache_x @@ -347,27 +336,17 @@ class WanResample(nn.Module): _feat_idx += 1 else: cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and _feat_cache[ - idx] is not None and _feat_cache[idx] != "Rep": - cache_x = torch.cat([ - _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) - if cache_x.shape[2] < 2 and _feat_cache[ - idx] is not None and _feat_cache[idx] == "Rep": - cache_x = torch.cat([ - torch.zeros_like(cache_x).to(cache_x.device), - cache_x - ], - dim=2) + if cache_x.shape[2] < 2 and _feat_cache[idx] is not None and _feat_cache[idx] != "Rep": + cache_x = torch.cat( + [_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + if cache_x.shape[2] < 2 and _feat_cache[idx] is not None and _feat_cache[idx] == "Rep": + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) x = self.time_conv(x) if _feat_cache[idx] == "Rep" else self.time_conv(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 x = x.reshape(b, 2, c, t, h, w) - x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), - 3) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) x = x.reshape(b, c, t * 2, h, w) feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) @@ -391,8 +370,7 @@ class WanResample(nn.Module): _feat_idx += 1 else: cache_x = x[:, :, -1:, :, :].clone() - x = self.time_conv( - torch.cat([_feat_cache[idx][:, :, -1:, :, :], x], 2)) + x = self.time_conv(torch.cat([_feat_cache[idx][:, :, -1:, :, :], x], 2)) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) @@ -431,8 +409,7 @@ class WanResidualBlock(nn.Module): self.norm2 = WanRMS_norm(out_dim, images=False) self.dropout = nn.Dropout(dropout) self.conv2 = WanCausalConv3d(out_dim, out_dim, 3, padding=1) - self.conv_shortcut = WanCausalConv3d( - in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + self.conv_shortcut = WanCausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() def forward(self, x): # Apply shortcut connection @@ -448,11 +425,7 @@ class WanResidualBlock(nn.Module): idx = _feat_idx cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: - cache_x = torch.cat([ - _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) + cache_x = torch.cat([_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv1(x, _feat_cache[idx]) _feat_cache[idx] = cache_x @@ -475,11 +448,7 @@ class WanResidualBlock(nn.Module): idx = _feat_idx cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: - cache_x = torch.cat([ - _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) + cache_x = torch.cat([_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv2(x, _feat_cache[idx]) _feat_cache[idx] = cache_x @@ -514,8 +483,7 @@ class WanAttentionBlock(nn.Module): identity = x batch_size, channels, time, height, width = x.size() - x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, - height, width) + x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width) x = self.norm(x) # compute query, key, value @@ -527,8 +495,7 @@ class WanAttentionBlock(nn.Module): # apply attention x = F.scaled_dot_product_attention(q, k, v) - x = x.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, - height, width) + x = x.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, height, width) # output projection x = self.proj(x) @@ -550,11 +517,7 @@ class WanMidBlock(nn.Module): non_linearity (str): Type of non-linearity to use. """ - def __init__(self, - dim: int, - dropout: float = 0.0, - non_linearity: str = "silu", - num_layers: int = 1): + def __init__(self, dim: int, dropout: float = 0.0, non_linearity: str = "silu", num_layers: int = 1): super().__init__() self.dim = dim @@ -582,7 +545,9 @@ class WanMidBlock(nn.Module): return x + class WanResidualDownBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout, num_res_blocks, temperal_downsample=False, down_flag=False): super().__init__() @@ -617,6 +582,7 @@ class WanResidualDownBlock(nn.Module): return x + self.avg_shortcut(x_copy) + class WanEncoder3d(nn.Module): r""" A 3D encoder module. @@ -633,17 +599,17 @@ class WanEncoder3d(nn.Module): """ def __init__( - self, - in_channels: int = 3, - dim=128, - z_dim=4, - dim_mult=(1, 2, 4, 4), - num_res_blocks=2, - attn_scales=(), - temperal_downsample=(True, True, False), - dropout=0.0, - non_linearity: str = "silu", - is_residual: bool = False, # wan 2.2 vae use a residual downblock + self, + in_channels: int = 3, + dim=128, + z_dim=4, + dim_mult=(1, 2, 4, 4), + num_res_blocks=2, + attn_scales=(), + temperal_downsample=(True, True, False), + dropout=0.0, + non_linearity: str = "silu", + is_residual: bool = False, # wan 2.2 vae use a residual downblock ): super().__init__() self.dim = dim @@ -664,8 +630,7 @@ class WanEncoder3d(nn.Module): # downsample blocks self.down_blocks = nn.ModuleList([]) - for i, (in_dim, - out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)): + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)): # residual (+attention) blocks if is_residual: self.down_blocks.append( @@ -676,8 +641,7 @@ class WanEncoder3d(nn.Module): num_res_blocks, temperal_downsample=temperal_downsample[i] if i != len(dim_mult) - 1 else False, down_flag=i != len(dim_mult) - 1, - ) - ) + )) else: for _ in range(num_res_blocks): self.down_blocks.append(WanResidualBlock(in_dim, out_dim, dropout)) @@ -692,10 +656,7 @@ class WanEncoder3d(nn.Module): scale /= 2.0 # middle blocks - self.mid_block = WanMidBlock(out_dim, - dropout, - non_linearity, - num_layers=1) + self.mid_block = WanMidBlock(out_dim, dropout, non_linearity, num_layers=1) # output blocks self.norm_out = WanRMS_norm(out_dim, images=False) @@ -711,11 +672,7 @@ class WanEncoder3d(nn.Module): cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: # cache last frame of last two chunk - cache_x = torch.cat([ - _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) + cache_x = torch.cat([_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv_in(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 @@ -742,11 +699,7 @@ class WanEncoder3d(nn.Module): cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: # cache last frame of last two chunk - cache_x = torch.cat([ - _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) + cache_x = torch.cat([_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv_out(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 @@ -756,6 +709,7 @@ class WanEncoder3d(nn.Module): x = self.conv_out(x) return x + # adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/autoencoder_kl_wan.py class WanResidualUpBlock(nn.Module): """ @@ -836,6 +790,7 @@ class WanResidualUpBlock(nn.Module): return x + class WanUpBlock(nn.Module): """ A block that handles upsampling for the WanVAE decoder. @@ -867,8 +822,7 @@ class WanUpBlock(nn.Module): # Add residual blocks and attention if needed current_dim = in_dim for _ in range(num_res_blocks + 1): - resnets.append( - WanResidualBlock(current_dim, out_dim, dropout, non_linearity)) + resnets.append(WanResidualBlock(current_dim, out_dim, dropout, non_linearity)) current_dim = out_dim self.resnets = nn.ModuleList(resnets) @@ -876,8 +830,7 @@ class WanUpBlock(nn.Module): # Add upsampling layer if needed self.upsamplers = None if upsample_mode is not None: - self.upsamplers = nn.ModuleList( - [WanResample(out_dim, mode=upsample_mode)]) + self.upsamplers = nn.ModuleList([WanResample(out_dim, mode=upsample_mode)]) self.gradient_checkpointing = False @@ -947,15 +900,11 @@ class WanDecoder3d(nn.Module): self.conv_in = WanCausalConv3d(z_dim, dims[0], 3, padding=1) # middle blocks - self.mid_block = WanMidBlock(dims[0], - dropout, - non_linearity, - num_layers=1) + self.mid_block = WanMidBlock(dims[0], dropout, non_linearity, num_layers=1) # upsample blocks self.up_blocks = nn.ModuleList([]) - for i, (in_dim, - out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)): + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)): # residual (+attention) blocks if i > 0 and not is_residual: # wan vae 2.1 @@ -1007,11 +956,7 @@ class WanDecoder3d(nn.Module): cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: # cache last frame of last two chunk - cache_x = torch.cat([ - _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) + cache_x = torch.cat([_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv_in(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 @@ -1037,11 +982,7 @@ class WanDecoder3d(nn.Module): cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: # cache last frame of last two chunk - cache_x = torch.cat([ - _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) + cache_x = torch.cat([_feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.conv_out(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 @@ -1051,13 +992,13 @@ class WanDecoder3d(nn.Module): x = self.conv_out(x) return x + def patchify(x, patch_size): if patch_size == 1: return x if x.dim() == 4: - x = rearrange( - x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) elif x.dim() == 5: x = rearrange( x, @@ -1070,13 +1011,13 @@ def patchify(x, patch_size): return x + def unpatchify(x, patch_size): if patch_size == 1: return x if x.dim() == 4: - x = rearrange( - x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) elif x.dim() == 5: x = rearrange( x, @@ -1096,8 +1037,8 @@ def _is_wan_vae_codec(name: str, submodule: nn.Module) -> bool: already-constructed DecodingStage. Without this, ``enable_torch_compile_vae`` falls back to replacing ``self.modules["vae"]`` and is a no-op for Wan. """ - return (name == "encoder" and isinstance(submodule, WanEncoder3d)) or ( - name == "decoder" and isinstance(submodule, WanDecoder3d)) + return (name == "encoder" and isinstance(submodule, WanEncoder3d)) or (name == "decoder" + and isinstance(submodule, WanDecoder3d)) class AutoencoderKLWan(nn.Module, ParallelTiledVAE): @@ -1129,9 +1070,9 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): if config.load_encoder: self.encoder = WanEncoder3d( in_channels=config.in_channels, - dim=config.base_dim, + dim=config.base_dim, z_dim=self.z_dim * 2, - dim_mult=config.dim_mult, + dim_mult=config.dim_mult, num_res_blocks=config.num_res_blocks, attn_scales=config.attn_scales, temperal_downsample=self.temperal_downsample, @@ -1143,12 +1084,12 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): if config.load_decoder: self.decoder = WanDecoder3d( - dim=decoder_base_dim, + dim=decoder_base_dim, z_dim=self.z_dim, - dim_mult=config.dim_mult, + dim_mult=config.dim_mult, num_res_blocks=config.num_res_blocks, attn_scales=config.attn_scales, - temperal_upsample=self.temperal_upsample, + temperal_upsample=self.temperal_upsample, dropout=config.dropout, out_channels=config.out_channels, is_residual=config.is_residual, @@ -1181,9 +1122,7 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): def _convert(module: nn.Module) -> None: for child in module.children(): if isinstance(child, nn.Conv3d): - child.weight.data = child.weight.data.to( - memory_format=torch.channels_last_3d - ) + child.weight.data = child.weight.data.to(memory_format=torch.channels_last_3d) else: _convert(child) @@ -1204,19 +1143,16 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): if i == 0: out = self.encoder(x[:, :, :1, :, :]) else: - out_ = self.encoder(x[:, :, - 1 + 4 * (i - 1):1 + 4 * i, :, :]) + out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :]) out = torch.cat([out, out_], 2) enc = self.quant_conv(out) - mu, logvar = enc[:, :self.z_dim, :, :, :], enc[:, - self.z_dim:, :, :, :] + mu, logvar = enc[:, :self.z_dim, :, :, :], enc[:, self.z_dim:, :, :, :] enc = torch.cat([mu, logvar], dim=1) enc = DiagonalGaussianDistribution(enc) self.clear_cache() else: for block in self.encoder.down_blocks: - if isinstance(block, - WanResample) and block.mode == "downsample3d": + if isinstance(block, WanResample) and block.mode == "downsample3d": _padding = list(block.time_conv._padding) _padding[4] = 2 block.time_conv._padding = tuple(_padding) @@ -1225,8 +1161,7 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): return enc def _encode(self, x: torch.Tensor, first_frame=False) -> torch.Tensor: - with forward_context(first_frame_arg=first_frame, - use_light_vae_arg=self.config.use_light_vae): + with forward_context(first_frame_arg=first_frame, use_light_vae_arg=self.config.use_light_vae): out = self.encoder(x) enc = self.quant_conv(out) mu, logvar = enc[:, :self.z_dim, :, :, :], enc[:, self.z_dim:, :, :, :] @@ -1271,7 +1206,7 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): if self.config.patch_size is not None: out = unpatchify(out, patch_size=self.config.patch_size) - + out = out.float() out = torch.clamp(out, min=-1.0, max=1.0) self.clear_cache() @@ -1282,8 +1217,7 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): def _decode(self, z: torch.Tensor, first_frame=False) -> torch.Tensor: x = self.post_quant_conv(z) - with forward_context(first_frame_arg=first_frame, - use_light_vae_arg=self.config.use_light_vae): + with forward_context(first_frame_arg=first_frame, use_light_vae_arg=self.config.use_light_vae): out = self.decoder(x) out = torch.clamp(out, min=-1.0, max=1.0) @@ -1311,6 +1245,7 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): return dec def get_streaming_cache(self) -> list[torch.Tensor | None]: + def _count_conv3d(model) -> int: count = 0 for m in model.modules(): @@ -1338,9 +1273,8 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): """ iter_ = z.shape[2] x = self.post_quant_conv(z) - - with forward_context(feat_cache_arg=cache, feat_idx_arg=0, - use_light_vae_arg=self.config.use_light_vae): + + with forward_context(feat_cache_arg=cache, feat_idx_arg=0, use_light_vae_arg=self.config.use_light_vae): outputs = [] for i in range(iter_): feat_idx.set(0) @@ -1375,5 +1309,6 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): dec = self.decode(z) return dec + # Entry point for model registry EntryClass = AutoencoderKLWan diff --git a/fastvideo/models/vision_utils.py b/fastvideo/models/vision_utils.py index 28aa016a1..feffe2178 100644 --- a/fastvideo/models/vision_utils.py +++ b/fastvideo/models/vision_utils.py @@ -22,8 +22,7 @@ from fastvideo.logger import init_logger logger = init_logger(__name__) -if version.parse(version.parse( - PIL.__version__).base_version) >= version.parse("9.1.0"): +if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"): PIL_INTERPOLATION = { "linear": PIL.Image.Resampling.BILINEAR, "bilinear": PIL.Image.Resampling.BILINEAR, @@ -95,9 +94,7 @@ def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: return 2.0 * images - 1.0 -def _fetch_image_bytes(url: str, - attempts: int = 3, - timeout: float = 30.0) -> bytes: +def _fetch_image_bytes(url: str, attempts: int = 3, timeout: float = 30.0) -> bytes: """Download ``url`` fully, retrying transient network errors. Reads the whole body via ``response.content`` instead of handing PIL a @@ -111,28 +108,23 @@ def _fetch_image_bytes(url: str, response.raise_for_status() return response.content except requests.RequestException as e: - status = getattr(getattr(e, "response", None), "status_code", - None) + status = getattr(getattr(e, "response", None), "status_code", None) # Permanent client errors (4xx except 408/429) won't heal on # retry — fail fast. - if (status is not None and 400 <= status < 500 - and status not in (408, 429)): + if (status is not None and 400 <= status < 500 and status not in (408, 429)): raise if attempt == attempts - 1: raise backoff = 2**attempt - logger.warning( - "Failed to download image %s (attempt %d/%d): %s. " - "Retrying in %ds.", url, attempt + 1, attempts, e, backoff) + logger.warning("Failed to download image %s (attempt %d/%d): %s. " + "Retrying in %ds.", url, attempt + 1, attempts, e, backoff) time.sleep(backoff) raise AssertionError("unreachable") # adapted from diffusers.utils import load_image -def load_image( - image: str | PIL.Image.Image, - convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None -) -> PIL.Image.Image: +def load_image(image: str | PIL.Image.Image, + convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None) -> PIL.Image.Image: """ Loads `image` to a PIL Image. @@ -160,8 +152,7 @@ def load_image( image = image else: raise ValueError( - "Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image." - ) + "Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image.") image = PIL.ImageOps.exif_transpose(image) @@ -202,8 +193,7 @@ def _load_gif(gif_path: str) -> tuple[list[PIL.Image.Image], float | None]: return pil_images, original_fps -def _load_video_with_ffmpeg( - video_path: str) -> tuple[list[PIL.Image.Image], float | None]: +def _load_video_with_ffmpeg(video_path: str) -> tuple[list[PIL.Image.Image], float | None]: """ Load frames from a video file using ffmpeg. @@ -220,9 +210,8 @@ def _load_video_with_ffmpeg( try: imageio.plugins.ffmpeg.get_exe() except AttributeError as e: - raise AttributeError( - "Unable to find an ffmpeg installation on your machine. " - "Please install via `uv pip install imageio-ffmpeg`") from e + raise AttributeError("Unable to find an ffmpeg installation on your machine. " + "Please install via `uv pip install imageio-ffmpeg`") from e pil_images = [] original_fps = None @@ -273,22 +262,18 @@ def load_video( if not (is_url or is_file): raise ValueError( - f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {video} is not a valid path." - ) + f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {video} is not a valid path.") if is_url: response = requests.get(video, stream=True) if response.status_code != 200: - raise ValueError( - f"Failed to download video. Status code: {response.status_code}" - ) + raise ValueError(f"Failed to download video. Status code: {response.status_code}") parsed_url = urlparse(video) file_name = os.path.basename(unquote(parsed_url.path)) suffix = os.path.splitext(file_name)[1] or ".mp4" - with tempfile.NamedTemporaryFile(suffix=suffix, - delete=False) as temp_file: + with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as temp_file: video_path = temp_file.name video_data = response.iter_content(chunk_size=8192) for chunk in video_data: @@ -357,8 +342,8 @@ def get_default_height_width( else: width = image.shape[2] - width, height = (x - x % vae_scale_factor for x in (width, height) - ) # resize to integer multiple of vae_scale_factor + width, height = (x - x % vae_scale_factor + for x in (width, height)) # resize to integer multiple of vae_scale_factor return height, width @@ -394,12 +379,10 @@ def resize( The resized image. """ if resize_mode != "default" and not isinstance(image, PIL.Image.Image): - raise ValueError( - f"Only PIL image input is supported for resize_mode {resize_mode}") + raise ValueError(f"Only PIL image input is supported for resize_mode {resize_mode}") assert isinstance(image, PIL.Image.Image) if resize_mode == "default": - image = image.resize((width, height), - resample=PIL_INTERPOLATION[resample]) + image = image.resize((width, height), resample=PIL_INTERPOLATION[resample]) elif resize_mode == "crop": src_w, src_h = image.size target_ratio = height / width @@ -415,14 +398,15 @@ def resize( top = max(0, (src_h - crop_h) // 2) left = max(0, (src_w - crop_w) // 2) image = image.crop((left, top, left + crop_w, top + crop_h)) - image = image.resize((width, height), - resample=PIL_INTERPOLATION[resample]) + image = image.resize((width, height), resample=PIL_INTERPOLATION[resample]) else: raise ValueError(f"resize_mode {resize_mode} is not supported") return image -def create_default_image(width: int = 512, height: int = 512, color: tuple[int, int, int] = (0, 0, 0)) -> PIL.Image.Image: +def create_default_image(width: int = 512, + height: int = 512, + color: tuple[int, int, int] = (0, 0, 0)) -> PIL.Image.Image: """ Create a default black PIL image. @@ -455,12 +439,8 @@ def preprocess_reference_image_for_clip(image: PIL.Image.Image, device: torch.de image_tensor = TF.to_tensor(image).sub_(0.5).div_(0.5).to(device) # Resize to CLIP's expected input size (224x224) using bicubic interpolation - resized_tensor = F.interpolate( - image_tensor.unsqueeze(0), - size=(224, 224), - mode='bicubic', - align_corners=False - ).squeeze(0) + resized_tensor = F.interpolate(image_tensor.unsqueeze(0), size=(224, 224), mode='bicubic', + align_corners=False).squeeze(0) # Denormalize back to [0, 1] range denormalized_tensor = resized_tensor.mul_(0.5).add_(0.5) diff --git a/fastvideo/tests/api/test_attention_selector_resolution.py b/fastvideo/tests/api/test_attention_selector_resolution.py index cc84c4916..00c813165 100644 --- a/fastvideo/tests/api/test_attention_selector_resolution.py +++ b/fastvideo/tests/api/test_attention_selector_resolution.py @@ -338,8 +338,8 @@ def test_component_with_no_concrete_decision_reports_no_request(monkeypatch): # (c) and it does NOT suppress the env var, which is the whole point monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") selector._cached_get_attn_backend.cache_clear() - passthrough = selector.get_attn_backend( - requested=selector.component_attention_backend(_Component(untouched)), **KWARGS) + passthrough = selector.get_attn_backend(requested=selector.component_attention_backend(_Component(untouched)), + **KWARGS) assert passthrough == selector.get_attn_backend(**KWARGS) diff --git a/fastvideo/tests/api/test_attn_qat_infer_capability_gate.py b/fastvideo/tests/api/test_attn_qat_infer_capability_gate.py index 4adc874ee..7094e8844 100644 --- a/fastvideo/tests/api/test_attn_qat_infer_capability_gate.py +++ b/fastvideo/tests/api/test_attn_qat_infer_capability_gate.py @@ -44,11 +44,7 @@ FALLBACK_CLASSES = { } -def _fake_gpu(monkeypatch, - *, - capability: tuple[int, int], - extension_imports: bool, - fa4_imports: bool = False) -> None: +def _fake_gpu(monkeypatch, *, capability: tuple[int, int], extension_imports: bool, fa4_imports: bool = False) -> None: monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: capability) monkeypatch.setattr( diff --git a/fastvideo/tests/api/test_cli_translation.py b/fastvideo/tests/api/test_cli_translation.py index a204daa7c..22bc7d033 100644 --- a/fastvideo/tests/api/test_cli_translation.py +++ b/fastvideo/tests/api/test_cli_translation.py @@ -59,9 +59,7 @@ def test_generate_parser_preserves_unknown_dotted_overrides(tmp_path): assert unknown == ["--request.sampling.seed", "42"] -def test_build_generate_run_config_loads_nested_config_and_dotted_overrides( - tmp_path, -): +def test_build_generate_run_config_loads_nested_config_and_dotted_overrides(tmp_path, ): config_path = tmp_path / "run.yaml" config_path.write_text( "generator:\n" @@ -188,8 +186,8 @@ def test_build_generate_run_config_rejects_flat_config(tmp_path): args, unknown = _parse_generate_args(["--config", str(config_path)]) with pytest.raises( - ValueError, - match="top-level 'generator' mapping", + ValueError, + match="top-level 'generator' mapping", ): build_generate_run_config(args, unknown) @@ -211,8 +209,8 @@ def test_build_generate_run_config_rejects_non_dotted_overrides(tmp_path): "2", ]) with pytest.raises( - ValueError, - match="CLI overrides must use dotted config paths", + ValueError, + match="CLI overrides must use dotted config paths", ): build_generate_run_config(args, unknown) @@ -226,8 +224,8 @@ def test_build_generate_run_config_requires_single_prompt_source(tmp_path): ) args, unknown = _parse_generate_args(["--config", str(missing_prompt_path)]) with pytest.raises( - ValueError, - match="Either request.prompt or request.inputs.prompt_path must be provided", + ValueError, + match="Either request.prompt or request.inputs.prompt_path must be provided", ): build_generate_run_config(args, unknown) @@ -243,8 +241,8 @@ def test_build_generate_run_config_requires_single_prompt_source(tmp_path): ) args, unknown = _parse_generate_args(["--config", str(conflicting_prompt_path)]) with pytest.raises( - ValueError, - match="Cannot provide both request.prompt and request.inputs.prompt_path", + ValueError, + match="Cannot provide both request.prompt and request.inputs.prompt_path", ): build_generate_run_config(args, unknown) @@ -287,8 +285,8 @@ def test_build_serve_config_rejects_flat_config(tmp_path): args, unknown = _parse_serve_args(["--config", str(config_path)]) with pytest.raises( - ValueError, - match="top-level 'generator' mapping", + ValueError, + match="top-level 'generator' mapping", ): build_serve_config(args, unknown) @@ -308,8 +306,8 @@ def test_build_serve_config_rejects_non_dotted_overrides(tmp_path): "9000", ]) with pytest.raises( - ValueError, - match="CLI overrides must use dotted config paths", + ValueError, + match="CLI overrides must use dotted config paths", ): build_serve_config(args, unknown) @@ -318,8 +316,8 @@ def test_generate_subcommand_requires_config(): args, _ = _parse_generate_args([]) with pytest.raises( - ValueError, - match="fastvideo generate requires --config PATH", + ValueError, + match="fastvideo generate requires --config PATH", ): GenerateSubcommand().validate(args) @@ -328,8 +326,8 @@ def test_serve_subcommand_requires_config(): args, _ = _parse_serve_args([]) with pytest.raises( - ValueError, - match="fastvideo serve requires --config PATH", + ValueError, + match="fastvideo serve requires --config PATH", ): ServeSubcommand().validate(args) @@ -351,8 +349,8 @@ def test_generate_subcommand_rejects_non_positive_num_gpus(tmp_path): ]) with pytest.raises( - ValueError, - match=r"generator\.engine\.num_gpus must be > 0; got 0", + ValueError, + match=r"generator\.engine\.num_gpus must be > 0; got 0", ): GenerateSubcommand().validate(args) @@ -372,8 +370,8 @@ def test_serve_subcommand_rejects_non_positive_num_gpus(tmp_path): ]) with pytest.raises( - ValueError, - match=r"generator\.engine\.num_gpus must be > 0; got 0", + ValueError, + match=r"generator\.engine\.num_gpus must be > 0; got 0", ): ServeSubcommand().validate(args) @@ -555,7 +553,7 @@ def test_streaming_run_server_rejects_missing_streaming_block(): streaming=None, ) with pytest.raises( - ValueError, - match="ServeConfig.streaming must be set", + ValueError, + match="ServeConfig.streaming must be set", ): streaming_server.run_server(config) diff --git a/fastvideo/tests/api/test_extra_overrides_routing.py b/fastvideo/tests/api/test_extra_overrides_routing.py index cf391fc26..0a181218b 100644 --- a/fastvideo/tests/api/test_extra_overrides_routing.py +++ b/fastvideo/tests/api/test_extra_overrides_routing.py @@ -24,8 +24,7 @@ import torch from fastvideo.api.sampling_param import SamplingParam from fastvideo.entrypoints.video_generator import ( - _BATCH_EXTRA_PASSTHROUGH_KEYS, -) + _BATCH_EXTRA_PASSTHROUGH_KEYS, ) from fastvideo.pipelines import ForwardBatch from fastvideo.utils import shallow_asdict @@ -48,9 +47,8 @@ def test_passthrough_keys_are_not_sampling_param_fields() -> None: import dataclasses sp_fields = {f.name for f in dataclasses.fields(SamplingParam())} leaked = sp_fields & set(_BATCH_EXTRA_PASSTHROUGH_KEYS) - assert not leaked, ( - f"Passthrough keys collide with SamplingParam fields: {leaked}. " - "Remove from _BATCH_EXTRA_PASSTHROUGH_KEYS or rename the field.") + assert not leaked, (f"Passthrough keys collide with SamplingParam fields: {leaked}. " + "Remove from _BATCH_EXTRA_PASSTHROUGH_KEYS or rename the field.") def test_sampling_param_update_rejects_unknown_keys() -> None: diff --git a/fastvideo/tests/api/test_ltx2_continuation.py b/fastvideo/tests/api/test_ltx2_continuation.py index 9bb9dd860..0ffca3473 100644 --- a/fastvideo/tests/api/test_ltx2_continuation.py +++ b/fastvideo/tests/api/test_ltx2_continuation.py @@ -38,9 +38,7 @@ from fastvideo.pipelines.basic.ltx2.continuation import ( def _make_typed_state() -> LTX2ContinuationState: return LTX2ContinuationState( segment_index=3, - video_frames=[ - (np.ones((64, 64, 3), dtype=np.uint8) * (i * 10)) for i in range(4) - ], + video_frames=[(np.ones((64, 64, 3), dtype=np.uint8) * (i * 10)) for i in range(4)], video_conditioning_frame_idx=9, video_conditioning_strength=0.75, audio_latents=torch.randn(1, 4, 16, 64, dtype=torch.float32), @@ -65,17 +63,12 @@ class TestRoundTrip: envelope = original.to_continuation_state() restored = LTX2ContinuationState.from_continuation_state(envelope) assert restored.segment_index == original.segment_index - assert restored.video_conditioning_frame_idx == ( - original.video_conditioning_frame_idx) - assert restored.video_conditioning_strength == ( - original.video_conditioning_strength) + assert restored.video_conditioning_frame_idx == (original.video_conditioning_frame_idx) + assert restored.video_conditioning_strength == (original.video_conditioning_strength) assert restored.audio_sample_rate == original.audio_sample_rate - assert restored.audio_conditioning_num_frames == ( - original.audio_conditioning_num_frames) - assert restored.audio_conditioning_strength == ( - original.audio_conditioning_strength) - assert restored.video_position_offset_sec == ( - original.video_position_offset_sec) + assert restored.audio_conditioning_num_frames == (original.audio_conditioning_num_frames) + assert restored.audio_conditioning_strength == (original.audio_conditioning_strength) + assert restored.video_position_offset_sec == (original.video_position_offset_sec) assert restored.metadata == original.metadata def test_inline_roundtrip_preserves_video_frames(self): @@ -84,8 +77,7 @@ class TestRoundTrip: restored = LTX2ContinuationState.from_continuation_state(envelope) assert restored.video_frames is not None assert len(restored.video_frames) == len(original.video_frames) - for before, after in zip(original.video_frames, - restored.video_frames): + for before, after in zip(original.video_frames, restored.video_frames): np.testing.assert_array_equal(before, after) def test_inline_roundtrip_preserves_audio_latents(self): @@ -93,11 +85,9 @@ class TestRoundTrip: envelope = original.to_continuation_state() restored = LTX2ContinuationState.from_continuation_state(envelope) assert restored.audio_latents is not None - assert tuple(restored.audio_latents.shape) == tuple( - original.audio_latents.shape) + assert tuple(restored.audio_latents.shape) == tuple(original.audio_latents.shape) assert restored.audio_latents.dtype == original.audio_latents.dtype - torch.testing.assert_close( - restored.audio_latents, original.audio_latents) + torch.testing.assert_close(restored.audio_latents, original.audio_latents) def test_payload_is_json_serializable(self): envelope = _make_typed_state().to_continuation_state() @@ -122,8 +112,7 @@ class TestRoundTrip: restored = LTX2ContinuationState.from_continuation_state(envelope) assert restored.audio_latents is not None assert restored.audio_latents.dtype == torch.bfloat16 - torch.testing.assert_close( - restored.audio_latents, state.audio_latents) + torch.testing.assert_close(restored.audio_latents, state.audio_latents) class TestBlobIndirection: @@ -149,12 +138,10 @@ class TestBlobIndirection: blob_store=blob_store, inline_threshold_bytes=0, ) - restored = LTX2ContinuationState.from_continuation_state( - envelope, blob_store=blob_store) + restored = LTX2ContinuationState.from_continuation_state(envelope, blob_store=blob_store) assert restored.video_frames is not None assert len(restored.video_frames) == len(original.video_frames) - torch.testing.assert_close( - restored.audio_latents, original.audio_latents) + torch.testing.assert_close(restored.audio_latents, original.audio_latents) def test_blob_id_held_when_store_unavailable(self): """Deserializing without a blob store preserves the blob id so @@ -197,14 +184,11 @@ class TestValidation: kind=LTX2_CONTINUATION_KIND, payload={"schema_version": 999}, ) - with pytest.raises(ValueError, - match="Unsupported LTX-2 continuation schema"): + with pytest.raises(ValueError, match="Unsupported LTX-2 continuation schema"): LTX2ContinuationState.from_continuation_state(envelope) def test_non_png_frame_rejected(self): - state = LTX2ContinuationState( - video_frames=[np.ones((64, 64, 3), dtype=np.float32)], - ) + state = LTX2ContinuationState(video_frames=[np.ones((64, 64, 3), dtype=np.float32)], ) with pytest.raises(ValueError, match="uint8 HxWx3"): state.to_continuation_state() @@ -232,14 +216,12 @@ class TestCompatLayerWireUp: def test_unknown_kind_rejected_at_boundary(self): from fastvideo.api.compat import _validate_continuation_state with pytest.raises(ValueError, match="Unknown ContinuationState kind"): - _validate_continuation_state( - ContinuationState(kind="mystery.v1", payload={})) + _validate_continuation_state(ContinuationState(kind="mystery.v1", payload={})) def test_empty_kind_rejected_at_boundary(self): from fastvideo.api.compat import _validate_continuation_state with pytest.raises(ValueError, match="non-empty string"): - _validate_continuation_state( - ContinuationState(kind="", payload={})) + _validate_continuation_state(ContinuationState(kind="", payload={})) def test_output_return_state_flag(self): request = GenerationRequest( diff --git a/fastvideo/tests/api/test_ltx2_gpu_pool_translation.py b/fastvideo/tests/api/test_ltx2_gpu_pool_translation.py index 610777f6e..77332e9d0 100644 --- a/fastvideo/tests/api/test_ltx2_gpu_pool_translation.py +++ b/fastvideo/tests/api/test_ltx2_gpu_pool_translation.py @@ -24,7 +24,6 @@ from fastvideo.api.compat import ( legacy_from_pretrained_to_config, ) - # Mirrors FastVideo-internal/ui/ltx2-streaming/server/gpu_pool.py # :lines 233-260 (load_kwargs constructed for VideoGenerator.from_pretrained). # @@ -108,8 +107,7 @@ class TestGpuPoolForwardTranslation: assert config.pipeline.components.config_root == "/models/ltx2-distilled/config" def test_refine_upsampler_routed_to_components(self, config) -> None: - assert config.pipeline.components.upsampler_weights == ( - "/models/ltx2-distilled/spatial_upsampler") + assert config.pipeline.components.upsampler_weights == ("/models/ltx2-distilled/spatial_upsampler") def test_empty_refine_lora_becomes_none(self, config) -> None: # gpu_pool passes "" to keep refine LoRA disabled; typed schema @@ -170,8 +168,7 @@ class TestGpuPoolReverseTranslation: assert args_kwargs["ltx2_refine_guidance_scale"] == 1.0 def test_refine_upsampler_path_reemitted(self, args_kwargs) -> None: - assert args_kwargs["ltx2_refine_upsampler_path"] == ( - "/models/ltx2-distilled/spatial_upsampler") + assert args_kwargs["ltx2_refine_upsampler_path"] == ("/models/ltx2-distilled/spatial_upsampler") def test_config_model_path_reemitted(self, args_kwargs) -> None: assert args_kwargs["config_model_path"] == "/models/ltx2-distilled/config" @@ -209,8 +206,7 @@ class TestRefineFlattenCoversAllTypedFields: def test_all_fields_reemitted(self, monkeypatch) -> None: from fastvideo import fastvideo_args as fva from fastvideo.api.compat import ( - generator_config_to_fastvideo_args, - ) + generator_config_to_fastvideo_args, ) from fastvideo.api.schema import GeneratorConfig, PipelineSelection from fastvideo.pipelines.basic.ltx2.stage_overrides import ( refine_preset_override_fields, @@ -240,10 +236,8 @@ class TestRefineFlattenCoversAllTypedFields: "image_crf": 18, "video_position_offset_sec": 2.5, } - all_fields = (refine_preset_override_fields() - | refine_stage_override_fields()) - assert set(refine_payload) == all_fields, ( - "payload must cover every typed field to exercise the flatten loop") + all_fields = (refine_preset_override_fields() | refine_stage_override_fields()) + assert set(refine_payload) == all_fields, ("payload must cover every typed field to exercise the flatten loop") config = GeneratorConfig( model_path="/models/ltx2", @@ -279,20 +273,25 @@ class TestCompileExtrasPreserved: kwargs = deepcopy(GPU_POOL_LOAD_KWARGS) kwargs["torch_compile_kwargs"] = { "backend": "inductor", - "options": {"triton.cudagraphs": False}, + "options": { + "triton.cudagraphs": False + }, "disable": False, } - config = legacy_from_pretrained_to_config( - "FastVideo/LTX2-Distilled-Diffusers", kwargs) + config = legacy_from_pretrained_to_config("FastVideo/LTX2-Distilled-Diffusers", kwargs) assert config.engine.compile.backend == "inductor" assert config.engine.compile.extras == { - "options": {"triton.cudagraphs": False}, + "options": { + "triton.cudagraphs": False + }, "disable": False, } generator_config_to_fastvideo_args(config) assert captured["torch_compile_kwargs"] == { "backend": "inductor", - "options": {"triton.cudagraphs": False}, + "options": { + "triton.cudagraphs": False + }, "disable": False, } diff --git a/fastvideo/tests/api/test_ltx2_param_mapping.py b/fastvideo/tests/api/test_ltx2_param_mapping.py index c48e46f74..b51ba4316 100644 --- a/fastvideo/tests/api/test_ltx2_param_mapping.py +++ b/fastvideo/tests/api/test_ltx2_param_mapping.py @@ -41,11 +41,9 @@ class TestLTX20ParamMappingDefault: ]) def test_vsa_gate_weight_not_renamed(self, source): target = _map(source, apply_gated_attention=False) - assert "to_gate_logits" not in target, ( - f"LTX-2.0 VSA gate {source!r} was rewritten to {target!r}; " - "expected pass-through with only prefix normalization.") - assert ".to_gate_compress." in target, ( - f"LTX-2.0 VSA gate {source!r} lost its identity in {target!r}.") + assert "to_gate_logits" not in target, (f"LTX-2.0 VSA gate {source!r} was rewritten to {target!r}; " + "expected pass-through with only prefix normalization.") + assert ".to_gate_compress." in target, (f"LTX-2.0 VSA gate {source!r} lost its identity in {target!r}.") @pytest.mark.parametrize("source", [ "transformer_blocks.0.attn1.to_gate_compress.lora_A.weight", @@ -56,43 +54,42 @@ class TestLTX20ParamMappingDefault: """``to_gate_compress`` is in ``DEFAULT_LORA_TARGET_MODULES``, so LoRAs trained with default targets ship these keys.""" target = _map(source, apply_gated_attention=False) - assert "to_gate_logits" not in target, ( - f"Default-target LoRA key {source!r} was rewritten to {target!r}; " - "this would break every LTX-2.0 LoRA in the wild.") + assert "to_gate_logits" not in target, (f"Default-target LoRA key {source!r} was rewritten to {target!r}; " + "this would break every LTX-2.0 LoRA in the wild.") def test_unrelated_param_still_prefix_stripped(self): """Generic prefix-strip behavior is unchanged in LTX-2.0 mode.""" assert (_map("diffusion_model.transformer_blocks.0.attn1.to_q.weight", - apply_gated_attention=False) == - "model.transformer_blocks.0.attn1.to_q.weight") + apply_gated_attention=False) == "model.transformer_blocks.0.attn1.to_q.weight") class TestLTX23ParamMappingGated: """``apply_gated_attention=True`` (LTX-2.3): the gated-attention ``to_gate_compress`` upstream key is renamed to ``to_gate_logits``.""" - @pytest.mark.parametrize("source,expected", [ - # No upstream prefix - ("transformer_blocks.0.attn1.to_gate_compress.weight", - "model.transformer_blocks.0.attn1.to_gate_logits.weight"), - # ``model.`` prefix - ("model.transformer_blocks.0.attn1.to_gate_compress.weight", - "model.transformer_blocks.0.attn1.to_gate_logits.weight"), - # ``diffusion_model.`` prefix - ("diffusion_model.transformer_blocks.0.attn1.to_gate_compress.weight", - "model.transformer_blocks.0.attn1.to_gate_logits.weight"), - # ``model.diffusion_model.`` prefix - ("model.diffusion_model.transformer_blocks.0.attn1.to_gate_compress.weight", - "model.transformer_blocks.0.attn1.to_gate_logits.weight"), - ]) + @pytest.mark.parametrize( + "source,expected", + [ + # No upstream prefix + ("transformer_blocks.0.attn1.to_gate_compress.weight", + "model.transformer_blocks.0.attn1.to_gate_logits.weight"), + # ``model.`` prefix + ("model.transformer_blocks.0.attn1.to_gate_compress.weight", + "model.transformer_blocks.0.attn1.to_gate_logits.weight"), + # ``diffusion_model.`` prefix + ("diffusion_model.transformer_blocks.0.attn1.to_gate_compress.weight", + "model.transformer_blocks.0.attn1.to_gate_logits.weight"), + # ``model.diffusion_model.`` prefix + ("model.diffusion_model.transformer_blocks.0.attn1.to_gate_compress.weight", + "model.transformer_blocks.0.attn1.to_gate_logits.weight"), + ]) def test_gate_rename_fires(self, source, expected): assert _map(source, apply_gated_attention=True) == expected def test_non_gate_param_still_prefix_stripped(self): """Non-gate params still get the prefix-strip behavior.""" assert (_map("diffusion_model.transformer_blocks.0.attn1.to_q.weight", - apply_gated_attention=True) == - "model.transformer_blocks.0.attn1.to_q.weight") + apply_gated_attention=True) == "model.transformer_blocks.0.attn1.to_q.weight") class TestParamMappingRuleOrdering: @@ -107,7 +104,6 @@ class TestParamMappingRuleOrdering: # prefix-normalized to ``model.transformer_blocks.0.attn1 # .to_gate_compress.weight``. first_pattern = next(iter(cfg.param_names_mapping)) - assert "to_gate_compress" in first_pattern, ( - "gate rename rule must be inserted at the front of " - "param_names_mapping; got first pattern: " - f"{first_pattern!r}") + assert "to_gate_compress" in first_pattern, ("gate rename rule must be inserted at the front of " + "param_names_mapping; got first pattern: " + f"{first_pattern!r}") diff --git a/fastvideo/tests/api/test_ltx2_stage_overrides.py b/fastvideo/tests/api/test_ltx2_stage_overrides.py index 2f207e7d7..44d08e5e1 100644 --- a/fastvideo/tests/api/test_ltx2_stage_overrides.py +++ b/fastvideo/tests/api/test_ltx2_stage_overrides.py @@ -99,8 +99,7 @@ class TestStageOverridesMirrorPresetSchema: def test_allowed_overrides_mirror_dataclass(self) -> None: import fastvideo.registry # noqa: F401 preset = get_preset("ltx2_two_stage", "ltx2") - refine_schema = next( - s for s in preset.stage_schemas if s.name == "refine") + refine_schema = next(s for s in preset.stage_schemas if s.name == "refine") assert refine_schema.allowed_overrides == refine_stage_override_fields() def test_roundtrip_through_validate_stage_overrides(self) -> None: @@ -110,12 +109,10 @@ class TestStageOverridesMirrorPresetSchema: num_inference_steps=3, guidance_scale=1.0, ) - validate_stage_overrides( - preset, {"refine": refine_override_to_dict(override)}) + validate_stage_overrides(preset, {"refine": refine_override_to_dict(override)}) def test_unknown_field_rejected(self) -> None: import fastvideo.registry # noqa: F401 preset = get_preset("ltx2_two_stage", "ltx2") with pytest.raises(ConfigValidationError): - validate_stage_overrides( - preset, {"refine": {"unknown_key": 1}}) + validate_stage_overrides(preset, {"refine": {"unknown_key": 1}}) diff --git a/fastvideo/tests/api/test_overrides.py b/fastvideo/tests/api/test_overrides.py index ddffd504d..1234d35a6 100644 --- a/fastvideo/tests/api/test_overrides.py +++ b/fastvideo/tests/api/test_overrides.py @@ -30,7 +30,11 @@ def test_parse_cli_overrides_casts_supported_scalar_and_collection_types() -> No "request.runtime.enable_teacache": True, "request.sampling.guidance_scale": 1.5, "request.prompt": ["a", "b"], - "request.extensions": {"ltx2": {"initial_latent_path": "/tmp/init.pt"}}, + "request.extensions": { + "ltx2": { + "initial_latent_path": "/tmp/init.pt" + } + }, "request.output.output_video_name": "clip", "request.state": None, } @@ -54,7 +58,9 @@ def test_apply_overrides_merges_nested_dicts_without_mutating_source() -> None: original = { "generator": { "model_path": "/models/base", - "engine": {"num_gpus": 1}, + "engine": { + "num_gpus": 1 + }, }, "request": {}, } @@ -74,20 +80,25 @@ def test_apply_overrides_merges_nested_dicts_without_mutating_source() -> None: def test_load_run_config_applies_dotted_overrides_before_validation(tmp_path) -> None: raw = { - "generator": {"model_path": "/models/base"}, - "request": {"prompt": "baseline"}, + "generator": { + "model_path": "/models/base" + }, + "request": { + "prompt": "baseline" + }, } path = tmp_path / "run.yaml" path.write_text(yaml.safe_dump(raw), encoding="utf-8") - loaded = load_run_config(path, overrides=[ - "--generator.pipeline.workload_type", - "t2v", - "--request.sampling.num_frames", - "81", - "--request.extensions.ltx2.initial_latent_path", - "/tmp/latent.pt", - ]) + loaded = load_run_config(path, + overrides=[ + "--generator.pipeline.workload_type", + "t2v", + "--request.sampling.num_frames", + "81", + "--request.extensions.ltx2.initial_latent_path", + "/tmp/latent.pt", + ]) assert loaded.generator.pipeline.workload_type == "t2v" assert loaded.request.sampling.num_frames == 81 diff --git a/fastvideo/tests/api/test_schema_parity_inventory.py b/fastvideo/tests/api/test_schema_parity_inventory.py index 8dc97e807..d95f79005 100644 --- a/fastvideo/tests/api/test_schema_parity_inventory.py +++ b/fastvideo/tests/api/test_schema_parity_inventory.py @@ -23,7 +23,6 @@ from fastvideo.entrypoints.openai.protocol import ( from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.utils import FlexibleArgumentParser - _REPO_ROOT = Path(__file__).resolve().parents[3] _INVENTORY_PATH = _REPO_ROOT / "docs" / "design" / "inference_schema_parity_inventory.yaml" @@ -69,19 +68,14 @@ def _get_extra_dataclass_fields( package = importlib.import_module(package_name) if not hasattr(package, "__path__"): continue - for _, modname, is_pkg in pkgutil.walk_packages( - package.__path__, prefix=f"{package_name}."): + for _, modname, is_pkg in pkgutil.walk_packages(package.__path__, prefix=f"{package_name}."): if modname.endswith(".__pycache__"): continue module = importlib.import_module(modname) for obj in vars(module).values(): - if (isinstance(obj, type) - and dataclasses.is_dataclass(obj) - and issubclass(obj, base_cls) + if (isinstance(obj, type) and dataclasses.is_dataclass(obj) and issubclass(obj, base_cls) and obj is not base_cls): - extras.update( - f.name for f in dataclasses.fields(obj) - if f.name not in base_fields) + extras.update(f.name for f in dataclasses.fields(obj) if f.name not in base_fields) return extras @@ -90,11 +84,7 @@ def _get_cli_dests(cmd_cls: type) -> set[str]: subparsers = parser.add_subparsers(dest="subparser") command = cmd_cls() subparser = command.subparser_init(subparsers) - return { - action.dest - for action in subparser._actions - if action.option_strings and action.dest != "help" - } + return {action.dest for action in subparser._actions if action.option_strings and action.dest != "help"} def _iter_inventory_targets(value: object) -> list[str]: @@ -137,8 +127,7 @@ def _walk_schema_target(root_type: type, target: str) -> None: return assert dataclasses.is_dataclass(current_annotation), ( f"{target!r} diverges at {'.'.join(target.split('.')[:depth - 1]) or ''}: " - f"{current_annotation!r} is not a dataclass or open dict boundary" - ) + f"{current_annotation!r} is not a dataclass or open dict boundary") hints = get_type_hints(current_annotation) assert segment in hints, f"{target!r} missing segment {segment!r}" current_annotation = hints[segment] @@ -316,8 +305,6 @@ def test_openai_seconds_mapping_preserves_duration_semantics( assert kwargs["fps"] == 24 assert kwargs["num_frames"] == 96 - seconds_entry = inventory["surfaces"]["openai_video_request"][ - "compatibility_only" - ]["seconds"] + seconds_entry = inventory["surfaces"]["openai_video_request"]["compatibility_only"]["seconds"] assert seconds_entry["target"] == "request.sampling.num_frames" assert "fps * seconds" in seconds_entry["note"] diff --git a/fastvideo/tests/api/test_typed_quant_flow.py b/fastvideo/tests/api/test_typed_quant_flow.py index 69e9ebeff..b7820d987 100644 --- a/fastvideo/tests/api/test_typed_quant_flow.py +++ b/fastvideo/tests/api/test_typed_quant_flow.py @@ -45,20 +45,17 @@ def captured_kwargs(monkeypatch): return captured -def test_typed_transformer_quant_resolves_to_nvfp4_instance( - captured_kwargs) -> None: +def test_typed_transformer_quant_resolves_to_nvfp4_instance(captured_kwargs) -> None: cfg = GeneratorConfig( model_path="FastVideo/LTX2-Distilled-Diffusers", - engine=EngineConfig( - quantization=QuantizationConfig(transformer_quant="NVFP4"), ), + engine=EngineConfig(quantization=QuantizationConfig(transformer_quant="NVFP4"), ), ) generator_config_to_fastvideo_args(cfg) - assert "transformer_quant" in captured_kwargs, ( - "compat layer must forward typed transformer_quant onto " - "FastVideoArgs.from_kwargs") - assert isinstance(captured_kwargs["transformer_quant"], NVFP4Config), ( - f"Expected NVFP4Config instance, got " - f"{type(captured_kwargs['transformer_quant']).__name__}") + assert "transformer_quant" in captured_kwargs, ("compat layer must forward typed transformer_quant onto " + "FastVideoArgs.from_kwargs") + assert isinstance(captured_kwargs["transformer_quant"], + NVFP4Config), (f"Expected NVFP4Config instance, got " + f"{type(captured_kwargs['transformer_quant']).__name__}") def test_no_typed_quant_omits_transformer_quant_kwarg(captured_kwargs) -> None: @@ -91,8 +88,7 @@ def test_apply_transformer_quant_pins_to_dit_config(monkeypatch) -> None: assert args.pipeline_config.dit_config.quant_config is nvfp4 -def test_apply_transformer_quant_does_not_overwrite_explicit_dit_config( -) -> None: +def test_apply_transformer_quant_does_not_overwrite_explicit_dit_config() -> None: """When the caller has explicitly set ``pipeline_config.dit_config.quant_config`` already, the typed carrier defers — the explicit setter wins. diff --git a/fastvideo/tests/api/test_validation_errors.py b/fastvideo/tests/api/test_validation_errors.py index 0525e1493..53873c535 100644 --- a/fastvideo/tests/api/test_validation_errors.py +++ b/fastvideo/tests/api/test_validation_errors.py @@ -6,15 +6,17 @@ from fastvideo.api import ConfigValidationError, RunConfig, parse_config def test_unknown_field_error_includes_nested_path() -> None: with pytest.raises( - ConfigValidationError, - match=r"generator\.engine\.bogus: unknown field", + ConfigValidationError, + match=r"generator\.engine\.bogus: unknown field", ): parse_config( RunConfig, { "generator": { "model_path": "/models/base", - "engine": {"bogus": True}, + "engine": { + "bogus": True + }, }, "request": {}, }, @@ -23,8 +25,8 @@ def test_unknown_field_error_includes_nested_path() -> None: def test_missing_required_field_error_includes_path() -> None: with pytest.raises( - ConfigValidationError, - match=r"generator\.model_path: missing required field", + ConfigValidationError, + match=r"generator\.model_path: missing required field", ): parse_config( RunConfig, @@ -37,15 +39,17 @@ def test_missing_required_field_error_includes_path() -> None: def test_invalid_literal_error_includes_path() -> None: with pytest.raises( - ConfigValidationError, - match=r"generator\.engine\.execution_backend: expected one of \['mp', 'ray'\]", + ConfigValidationError, + match=r"generator\.engine\.execution_backend: expected one of \['mp', 'ray'\]", ): parse_config( RunConfig, { "generator": { "model_path": "/models/base", - "engine": {"execution_backend": "threaded"}, + "engine": { + "execution_backend": "threaded" + }, }, "request": {}, }, @@ -54,21 +58,21 @@ def test_invalid_literal_error_includes_path() -> None: def test_invalid_nested_type_error_includes_list_path() -> None: with pytest.raises( - ConfigValidationError, - match=r"request\.plan\.stages\[0\]\.name: expected str", + ConfigValidationError, + match=r"request\.plan\.stages\[0\]\.name: expected str", ): parse_config( RunConfig, { - "generator": {"model_path": "/models/base"}, + "generator": { + "model_path": "/models/base" + }, "request": { "plan": { - "stages": [ - { - "name": 123, - "kind": "sample", - } - ] + "stages": [{ + "name": 123, + "kind": "sample", + }] } }, }, diff --git a/fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py b/fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py index 880c5dce1..e579ddac5 100644 --- a/fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py +++ b/fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py @@ -144,11 +144,9 @@ def _register_fa4_quantize_cpu_kernel(): def _cpu_kernel(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: batch, seqlen, nheads, headdim = tensor_4d.shape seqlen_padded = (seqlen + 127) // 128 * 128 - fp4 = torch.zeros(batch, seqlen_padded, nheads, headdim // 2, - dtype=torch.int8).view(torch.float4_e2m1fn_x2) + fp4 = torch.zeros(batch, seqlen_padded, nheads, headdim // 2, dtype=torch.int8).view(torch.float4_e2m1fn_x2) rest_m, rest_k = seqlen_padded // 128, (headdim // 16) // 4 - sf = torch.zeros(batch, nheads, rest_m, rest_k, 32, 4, 4, - dtype=torch.uint8).permute(4, 5, 2, 6, 3, 1, 0) + sf = torch.zeros(batch, nheads, rest_m, rest_k, 32, 4, 4, dtype=torch.uint8).permute(4, 5, 2, 6, 3, 1, 0) return fp4, sf try: @@ -197,6 +195,6 @@ def test_fa4_quantize_op_fake_matches_real() -> None: x = torch.randn(2, 200, 2, 128, dtype=torch.bfloat16) torch.library.opcheck( torch.ops.fastvideo.nvfp4_quantize_fa4, - (x,), + (x, ), test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"), ) diff --git a/fastvideo/tests/attention/test_flash_attn_cute_custom_op.py b/fastvideo/tests/attention/test_flash_attn_cute_custom_op.py index 2fc540937..a1d0a5167 100644 --- a/fastvideo/tests/attention/test_flash_attn_cute_custom_op.py +++ b/fastvideo/tests/attention/test_flash_attn_cute_custom_op.py @@ -10,7 +10,6 @@ from types import SimpleNamespace import pytest import torch - _MODULE_NAME = "_fastvideo_flash_attn_cute_for_tests" @@ -18,12 +17,7 @@ def _load_local_flash_attn_cute_module(): if _MODULE_NAME in sys.modules: return sys.modules[_MODULE_NAME] - module_path = ( - Path(__file__).resolve().parents[2] - / "attention" - / "utils" - / "flash_attn_cute.py" - ) + module_path = (Path(__file__).resolve().parents[2] / "attention" / "utils" / "flash_attn_cute.py") spec = importlib.util.spec_from_file_location(_MODULE_NAME, module_path) if spec is None or spec.loader is None: raise RuntimeError(f"Failed to load module spec from {module_path}") @@ -101,9 +95,7 @@ def flash_attn_impls(): major, _ = torch.cuda.get_device_capability() if major not in (9, 10, 11): - pytest.skip( - f"flash_attn.cute supports SM 9/10/11, got compute capability {major}" - ) + pytest.skip(f"flash_attn.cute supports SM 9/10/11, got compute capability {major}") try: from flash_attn.cute import ( @@ -124,9 +116,7 @@ def flash_attn_impls(): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("causal", [False, True]) -def test_flash_attn_func_parity_forward_backward( - flash_attn_impls, dtype: torch.dtype, causal: bool -): +def test_flash_attn_func_parity_forward_backward(flash_attn_impls, dtype: torch.dtype, causal: bool): if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported(): pytest.skip("bfloat16 is not supported on this GPU") @@ -152,33 +142,31 @@ def test_flash_attn_func_parity_forward_backward( k_test = k_ref.detach().clone().requires_grad_(True) v_test = v_ref.detach().clone().requires_grad_(True) - out_ref = _extract_out(ref_flash_attn_func( - q_ref, - k_ref, - v_ref, - softmax_scale=None, - causal=causal, - deterministic=False, - )) - out_test = _extract_out(custom_flash_attn_func( - q_test, - k_test, - v_test, - dropout_p=0.0, - softmax_scale=None, - causal=causal, - deterministic=False, - )) + out_ref = _extract_out( + ref_flash_attn_func( + q_ref, + k_ref, + v_ref, + softmax_scale=None, + causal=causal, + deterministic=False, + )) + out_test = _extract_out( + custom_flash_attn_func( + q_test, + k_test, + v_test, + dropout_p=0.0, + softmax_scale=None, + causal=causal, + deterministic=False, + )) _assert_close(out_test, out_ref, dtype=dtype) dout = torch.randn_like(out_ref) - dq_ref, dk_ref, dv_ref = torch.autograd.grad( - (out_ref * dout).sum(), (q_ref, k_ref, v_ref) - ) - dq_test, dk_test, dv_test = torch.autograd.grad( - (out_test * dout).sum(), (q_test, k_test, v_test) - ) + dq_ref, dk_ref, dv_ref = torch.autograd.grad((out_ref * dout).sum(), (q_ref, k_ref, v_ref)) + dq_test, dk_test, dv_test = torch.autograd.grad((out_test * dout).sum(), (q_test, k_test, v_test)) _assert_close(dq_test, dq_ref, dtype=dtype, is_grad=True) _assert_close(dk_test, dk_ref, dtype=dtype, is_grad=True) @@ -187,9 +175,7 @@ def test_flash_attn_func_parity_forward_backward( @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("causal", [False, True]) -def test_flash_attn_varlen_func_parity_forward_backward( - flash_attn_impls, dtype: torch.dtype, causal: bool -): +def test_flash_attn_varlen_func_parity_forward_backward(flash_attn_impls, dtype: torch.dtype, causal: bool): if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported(): pytest.skip("bfloat16 is not supported on this GPU") @@ -211,9 +197,7 @@ def test_flash_attn_varlen_func_parity_forward_backward( total = int(cu[-1].item()) max_seqlen = max(seqlens) - q_ref = torch.randn( - total, nheads, headdim, device=device, dtype=dtype, requires_grad=True - ) + q_ref = torch.randn(total, nheads, headdim, device=device, dtype=dtype, requires_grad=True) k_ref = torch.randn_like(q_ref, requires_grad=True) v_ref = torch.randn_like(q_ref, requires_grad=True) @@ -221,41 +205,39 @@ def test_flash_attn_varlen_func_parity_forward_backward( k_test = k_ref.detach().clone().requires_grad_(True) v_test = v_ref.detach().clone().requires_grad_(True) - out_ref = _extract_out(ref_flash_attn_varlen_func( - q_ref, - k_ref, - v_ref, - cu_seqlens_q=cu, - cu_seqlens_k=cu, - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, - softmax_scale=None, - causal=causal, - deterministic=False, - )) - out_test = _extract_out(custom_flash_attn_varlen_func( - q_test, - k_test, - v_test, - cu, - cu, - max_seqlen, - max_seqlen, - dropout_p=0.0, - softmax_scale=None, - causal=causal, - deterministic=False, - )) + out_ref = _extract_out( + ref_flash_attn_varlen_func( + q_ref, + k_ref, + v_ref, + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + softmax_scale=None, + causal=causal, + deterministic=False, + )) + out_test = _extract_out( + custom_flash_attn_varlen_func( + q_test, + k_test, + v_test, + cu, + cu, + max_seqlen, + max_seqlen, + dropout_p=0.0, + softmax_scale=None, + causal=causal, + deterministic=False, + )) _assert_close(out_test, out_ref, dtype=dtype) dout = torch.randn_like(out_ref) - dq_ref, dk_ref, dv_ref = torch.autograd.grad( - (out_ref * dout).sum(), (q_ref, k_ref, v_ref) - ) - dq_test, dk_test, dv_test = torch.autograd.grad( - (out_test * dout).sum(), (q_test, k_test, v_test) - ) + dq_ref, dk_ref, dv_ref = torch.autograd.grad((out_ref * dout).sum(), (q_ref, k_ref, v_ref)) + dq_test, dk_test, dv_test = torch.autograd.grad((out_test * dout).sum(), (q_test, k_test, v_test)) _assert_close(dq_test, dq_ref, dtype=dtype, is_grad=True) _assert_close(dk_test, dk_ref, dtype=dtype, is_grad=True) diff --git a/fastvideo/tests/attention/test_flash_attn_default_custom_op.py b/fastvideo/tests/attention/test_flash_attn_default_custom_op.py index eb58ff011..4c25a1695 100644 --- a/fastvideo/tests/attention/test_flash_attn_default_custom_op.py +++ b/fastvideo/tests/attention/test_flash_attn_default_custom_op.py @@ -41,10 +41,8 @@ def fa_default_impls(): pytest.skip(f"flash_attn_default not importable: {exc}") if fa_module.fa_version not in ("2", "3"): - pytest.skip( - f"FA2/FA3 default custom op only exists for fa_version in (2, 3); " - f"got {fa_module.fa_version!r}" - ) + pytest.skip(f"FA2/FA3 default custom op only exists for fa_version in (2, 3); " + f"got {fa_module.fa_version!r}") # compilable dispatcher, the original FA wrapper it falls back to, the # raw custom op for opcheck, and the fa_version (FA2 has full register_ @@ -60,8 +58,7 @@ def fa_default_impls(): def _qkv(dtype, requires_grad): device = torch.device("cuda") batch, seqlen, nheads, headdim = 2, 64, 4, 64 - q = torch.randn(batch, seqlen, nheads, headdim, device=device, - dtype=dtype, requires_grad=requires_grad) + q = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype, requires_grad=requires_grad) k = torch.randn_like(q, requires_grad=requires_grad) v = torch.randn_like(q, requires_grad=requires_grad) return q, k, v @@ -107,10 +104,8 @@ def test_default_compilable_training_backward_flows(fa_default_impls, dtype, cau out_test = compilable(q_test, k_test, v_test, softmax_scale=None, causal=causal) dout = torch.randn_like(out_ref) - dq_ref, dk_ref, dv_ref = torch.autograd.grad( - (out_ref * dout).sum(), (q_ref, k_ref, v_ref)) - dq_test, dk_test, dv_test = torch.autograd.grad( - (out_test * dout).sum(), (q_test, k_test, v_test)) + dq_ref, dk_ref, dv_ref = torch.autograd.grad((out_ref * dout).sum(), (q_ref, k_ref, v_ref)) + dq_test, dk_test, dv_test = torch.autograd.grad((out_test * dout).sum(), (q_test, k_test, v_test)) atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2 torch.testing.assert_close(dq_test, dq_ref, atol=atol, rtol=rtol) @@ -156,15 +151,14 @@ def test_default_op_backward_through_registered_autograd(fa_default_impls, dtype # Custom-op grads via the registered backward — unpack (out, lse), discard lse. out_test, _ = op(q_test, k_test, v_test, None, causal) - torch.testing.assert_close(out_test, out_ref, + torch.testing.assert_close(out_test, + out_ref, atol=0 if dtype == torch.float16 else 1e-3, rtol=0 if dtype == torch.float16 else 1e-3) dout = torch.randn_like(out_ref) - dq_ref, dk_ref, dv_ref = torch.autograd.grad( - (out_ref * dout).sum(), (q_ref, k_ref, v_ref)) - dq_test, dk_test, dv_test = torch.autograd.grad( - (out_test * dout).sum(), (q_test, k_test, v_test)) + dq_ref, dk_ref, dv_ref = torch.autograd.grad((out_ref * dout).sum(), (q_ref, k_ref, v_ref)) + dq_test, dk_test, dv_test = torch.autograd.grad((out_test * dout).sum(), (q_test, k_test, v_test)) atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2 torch.testing.assert_close(dq_test, dq_ref, atol=atol, rtol=rtol) diff --git a/fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py b/fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py index fb0f7300a..a267693d4 100644 --- a/fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py +++ b/fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py @@ -52,10 +52,8 @@ def _dtype_skip(dtype): def _fa2_only(mod): if mod._FA_VARLEN_VERSION != "2": - pytest.skip( - f"register_autograd is only wired for FA2; got " - f"_FA_VARLEN_VERSION={mod._FA_VARLEN_VERSION!r}" - ) + pytest.skip(f"register_autograd is only wired for FA2; got " + f"_FA_VARLEN_VERSION={mod._FA_VARLEN_VERSION!r}") def _padding_mask(batch, seqlen, valid_lens, device): @@ -107,8 +105,8 @@ def test_no_pad_training_backward_through_registered_autograd(no_pad_impls, dtyp out_test = mod.flash_attn_no_pad_compilable(qkv_test, mask, causal=False, dropout_p=0.0) dout = torch.randn_like(out_ref) - (dqkv_ref,) = torch.autograd.grad((out_ref * dout).sum(), (qkv_ref,)) - (dqkv_test,) = torch.autograd.grad((out_test * dout).sum(), (qkv_test,)) + (dqkv_ref, ) = torch.autograd.grad((out_ref * dout).sum(), (qkv_ref, )) + (dqkv_test, ) = torch.autograd.grad((out_test * dout).sum(), (qkv_test, )) atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2 torch.testing.assert_close(dqkv_test, dqkv_ref, atol=atol, rtol=rtol) @@ -191,8 +189,13 @@ def test_varlen_qk_training_backward_through_registered_autograd(no_pad_impls, d v_test = v_ref.detach().clone().requires_grad_(True) out_ref = mod.flash_attn_varlen_qk_no_pad(q_ref, k_ref, v_ref, qmask, kmask, causal=False, dropout_p=0.0) - out_test = mod.flash_attn_varlen_qk_no_pad_compilable(q_test, k_test, v_test, qmask, kmask, - causal=False, dropout_p=0.0) + out_test = mod.flash_attn_varlen_qk_no_pad_compilable(q_test, + k_test, + v_test, + qmask, + kmask, + causal=False, + dropout_p=0.0) dout = torch.randn_like(out_ref) dq_ref, dk_ref, dv_ref = torch.autograd.grad((out_ref * dout).sum(), (q_ref, k_ref, v_ref)) diff --git a/fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py b/fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py index f0a12e50e..a49883fb7 100644 --- a/fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py +++ b/fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py @@ -26,8 +26,9 @@ import sys import pytest if importlib.util.find_spec("flash_attn") is None: - pytest.skip("flash_attn not installed; resolver tests require it for the unconditional top-level imports in flash_attn_no_pad", - allow_module_level=True) + pytest.skip( + "flash_attn not installed; resolver tests require it for the unconditional top-level imports in flash_attn_no_pad", + allow_module_level=True) def _reload_resolver_module(): diff --git a/fastvideo/tests/attention/test_selector_role_override.py b/fastvideo/tests/attention/test_selector_role_override.py index 602386114..d894f3719 100644 --- a/fastvideo/tests/attention/test_selector_role_override.py +++ b/fastvideo/tests/attention/test_selector_role_override.py @@ -41,8 +41,7 @@ def test_role_override_invalidates_cached_backend(monkeypatch) -> None: try: assert selector.get_attn_backend(**kwargs) == "FLASH_ATTN" - with selector._component_attention_backend_scope( - AttentionBackendEnum.TORCH_SDPA, component="role"): + with selector._component_attention_backend_scope(AttentionBackendEnum.TORCH_SDPA, component="role"): # Same cache key as above; the role-local request must still win. assert selector.get_attn_backend(**kwargs) == "TORCH_SDPA" @@ -59,8 +58,7 @@ def test_unsupported_role_override_honors_layer_default(monkeypatch) -> None: monkeypatch.setattr(selector, "resolve_obj_by_qualname", lambda name: name) try: - with selector._component_attention_backend_scope( - AttentionBackendEnum.ATTN_QAT_TRAIN, component="role"): + with selector._component_attention_backend_scope(AttentionBackendEnum.ATTN_QAT_TRAIN, component="role"): assert selector.get_attn_backend( head_size=128, dtype=torch.bfloat16, diff --git a/fastvideo/tests/audio/test_ltx2_bwe.py b/fastvideo/tests/audio/test_ltx2_bwe.py index 773174e4d..e23618405 100644 --- a/fastvideo/tests/audio/test_ltx2_bwe.py +++ b/fastvideo/tests/audio/test_ltx2_bwe.py @@ -26,7 +26,6 @@ from fastvideo.models.audio.ltx2_audio_vae import ( _build_stft_basis, ) - _FILTER_LENGTH = 512 _HOP_LENGTH = 80 _WIN_LENGTH = 512 @@ -39,9 +38,8 @@ class TestSTFTBasisDeterministic: def test_forward_basis_is_not_zero(self): stft = _STFTFn(_FILTER_LENGTH, _HOP_LENGTH, _WIN_LENGTH) - assert torch.any(stft.forward_basis != 0), ( - "forward_basis is all-zero — the BWE STFT will silently produce " - "zero magnitudes (regression of #1398's S1 bug).") + assert torch.any(stft.forward_basis != 0), ("forward_basis is all-zero — the BWE STFT will silently produce " + "zero magnitudes (regression of #1398's S1 bug).") def test_forward_basis_shape(self): stft = _STFTFn(_FILTER_LENGTH, _HOP_LENGTH, _WIN_LENGTH) @@ -56,9 +54,8 @@ class TestSTFTBasisDeterministic: t = torch.arange(n, dtype=torch.float32) / _SAMPLING_RATE wave = torch.sin(2 * math.pi * 1000.0 * t).unsqueeze(0) # (1, n) magnitude, phase = stft(wave) - assert torch.any(magnitude != 0), ( - "STFT magnitude was zero for a non-zero waveform — basis init " - "regressed.") + assert torch.any(magnitude != 0), ("STFT magnitude was zero for a non-zero waveform — basis init " + "regressed.") assert magnitude.shape == phase.shape assert magnitude.shape[1] == _FILTER_LENGTH // 2 + 1 @@ -67,44 +64,23 @@ class TestMelBasisDeterministic: """``MelSTFT`` no longer leaves ``mel_basis`` as zeros.""" def test_mel_basis_is_not_zero(self): - mel = MelSTFT(_FILTER_LENGTH, - _HOP_LENGTH, - _WIN_LENGTH, - _N_MELS, - sampling_rate=_SAMPLING_RATE) - assert torch.any(mel.mel_basis != 0), ( - "mel_basis is all-zero — log-mel will collapse to log(1e-5) " - "(regression of #1398's S1 bug).") + mel = MelSTFT(_FILTER_LENGTH, _HOP_LENGTH, _WIN_LENGTH, _N_MELS, sampling_rate=_SAMPLING_RATE) + assert torch.any(mel.mel_basis != 0), ("mel_basis is all-zero — log-mel will collapse to log(1e-5) " + "(regression of #1398's S1 bug).") def test_mel_basis_shape(self): - mel = MelSTFT(_FILTER_LENGTH, - _HOP_LENGTH, - _WIN_LENGTH, - _N_MELS, - sampling_rate=_SAMPLING_RATE) + mel = MelSTFT(_FILTER_LENGTH, _HOP_LENGTH, _WIN_LENGTH, _N_MELS, sampling_rate=_SAMPLING_RATE) assert mel.mel_basis.shape == (_N_MELS, _FILTER_LENGTH // 2 + 1) def test_mel_basis_is_non_negative(self): """Triangular mel filters are non-negative everywhere.""" - mel = MelSTFT(_FILTER_LENGTH, - _HOP_LENGTH, - _WIN_LENGTH, - _N_MELS, - sampling_rate=_SAMPLING_RATE) + mel = MelSTFT(_FILTER_LENGTH, _HOP_LENGTH, _WIN_LENGTH, _N_MELS, sampling_rate=_SAMPLING_RATE) assert torch.all(mel.mel_basis >= 0) def test_mel_basis_changes_with_sampling_rate(self): """Different sample rates -> different mel filterbanks.""" - mel_16k = MelSTFT(_FILTER_LENGTH, - _HOP_LENGTH, - _WIN_LENGTH, - _N_MELS, - sampling_rate=16000) - mel_48k = MelSTFT(_FILTER_LENGTH, - _HOP_LENGTH, - _WIN_LENGTH, - _N_MELS, - sampling_rate=48000) + mel_16k = MelSTFT(_FILTER_LENGTH, _HOP_LENGTH, _WIN_LENGTH, _N_MELS, sampling_rate=16000) + mel_48k = MelSTFT(_FILTER_LENGTH, _HOP_LENGTH, _WIN_LENGTH, _N_MELS, sampling_rate=48000) assert not torch.allclose(mel_16k.mel_basis, mel_48k.mel_basis) @@ -112,11 +88,7 @@ class TestMelSpectrogramSanity: """End-to-end sanity: a 1 kHz tone produces a peaked log-mel, not constant.""" def _build(self) -> MelSTFT: - return MelSTFT(_FILTER_LENGTH, - _HOP_LENGTH, - _WIN_LENGTH, - _N_MELS, - sampling_rate=_SAMPLING_RATE) + return MelSTFT(_FILTER_LENGTH, _HOP_LENGTH, _WIN_LENGTH, _N_MELS, sampling_rate=_SAMPLING_RATE) def test_log_mel_is_not_constant(self): mel = self._build() @@ -126,9 +98,8 @@ class TestMelSpectrogramSanity: log_mel, magnitude, phase, energy = mel.mel_spectrogram(wave) # If forward_basis or mel_basis were zero, log_mel would collapse # to log(1e-5) ≈ -11.51 (a constant). - assert log_mel.std().item() > 1e-3, ( - "log-mel is effectively constant; the deterministic basis init " - "did not take effect (regression of #1398's S1 bug).") + assert log_mel.std().item() > 1e-3, ("log-mel is effectively constant; the deterministic basis init " + "did not take effect (regression of #1398's S1 bug).") def test_log_mel_shapes_match_input(self): mel = self._build() @@ -149,8 +120,7 @@ class TestStandaloneBasisHelpers: def test_stft_basis_helper_nonzero(self): basis = _build_stft_basis(_FILTER_LENGTH, _WIN_LENGTH) assert torch.any(basis != 0) - assert basis.shape == (2 * (_FILTER_LENGTH // 2 + 1), 1, - _FILTER_LENGTH) + assert basis.shape == (2 * (_FILTER_LENGTH // 2 + 1), 1, _FILTER_LENGTH) assert basis.dtype == torch.float32 def test_mel_basis_helper_nonzero(self): diff --git a/fastvideo/tests/conftest.py b/fastvideo/tests/conftest.py index 35657b7c8..66c1e71ce 100644 --- a/fastvideo/tests/conftest.py +++ b/fastvideo/tests/conftest.py @@ -4,8 +4,7 @@ import pytest import torch import numpy as np -from fastvideo.distributed import (maybe_init_distributed_environment_and_model_parallel, - cleanup_dist_env_and_memory) +from fastvideo.distributed import (maybe_init_distributed_environment_and_model_parallel, cleanup_dist_env_and_memory) @pytest.fixture(scope="function") diff --git a/fastvideo/tests/contract/test_ci_test_collection.py b/fastvideo/tests/contract/test_ci_test_collection.py index 8150579b1..9b0152db2 100644 --- a/fastvideo/tests/contract/test_ci_test_collection.py +++ b/fastvideo/tests/contract/test_ci_test_collection.py @@ -48,37 +48,27 @@ def _dirs_with_tests() -> list[str]: def _ci_text() -> str: - return "\n".join( - src.read_text(errors="replace") for src in CI_SOURCES if src.exists()) + return "\n".join(src.read_text(errors="replace") for src in CI_SOURCES if src.exists()) def test_every_test_directory_is_collected_or_allowlisted(): ci_text = _ci_text() - dark = [ - name for name in _dirs_with_tests() - if f"tests/{name}" not in ci_text and name not in ALLOWLIST - ] - assert not dark, ( - f"Test directories not referenced by any CI lane and not " - f"allowlisted: {dark}. Wire them into a lane in " - f"fastvideo/tests/modal/pr_test.py (or a Buildkite step), or add an " - f"allowlist entry with a reason in {__file__}.") + dark = [name for name in _dirs_with_tests() if f"tests/{name}" not in ci_text and name not in ALLOWLIST] + assert not dark, (f"Test directories not referenced by any CI lane and not " + f"allowlisted: {dark}. Wire them into a lane in " + f"fastvideo/tests/modal/pr_test.py (or a Buildkite step), or add an " + f"allowlist entry with a reason in {__file__}.") def test_local_tests_stays_out_of_ci(): # tests/local_tests/ (repo root) is developer-local by design (author # decision, 2026-07-05): parity scaffolds and machine-specific checks # that must never gate CI. Fail if any CI source starts collecting it. - assert "tests/local_tests" not in _ci_text(), ( - "tests/local_tests/ is local-only by design; remove the CI " - "reference or move the tests into a fastvideo/tests/ lane.") + assert "tests/local_tests" not in _ci_text(), ("tests/local_tests/ is local-only by design; remove the CI " + "reference or move the tests into a fastvideo/tests/ lane.") def test_allowlist_entries_are_still_real_directories(): # A stale allowlist hides regressions; entries must track reality. - missing = [ - name for name in ALLOWLIST - if name != "modal" and not (TESTS_ROOT / name).is_dir() - ] - assert not missing, ( - f"Allowlisted directories no longer exist — remove them: {missing}") + missing = [name for name in ALLOWLIST if name != "modal" and not (TESTS_ROOT / name).is_dir()] + assert not missing, (f"Allowlisted directories no longer exist — remove them: {missing}") diff --git a/fastvideo/tests/contract/test_dreamverse_shape.py b/fastvideo/tests/contract/test_dreamverse_shape.py index b418b09a5..0def46114 100644 --- a/fastvideo/tests/contract/test_dreamverse_shape.py +++ b/fastvideo/tests/contract/test_dreamverse_shape.py @@ -72,17 +72,14 @@ class TestDreamverseLoadKwargsShape: field, not in the ``experimental`` escape hatch.""" def test_all_kwargs_land_on_typed_fields(self): - config = legacy_from_pretrained_to_config( - "/models/ltx2", _dreamverse_load_kwargs()) + config = legacy_from_pretrained_to_config("/models/ltx2", _dreamverse_load_kwargs()) assert isinstance(config, GeneratorConfig) # None of the kwargs should have been routed to experimental. - assert config.pipeline.experimental == {}, ( - "Dreamverse kwargs leaked into pipeline.experimental: " - f"{config.pipeline.experimental}") + assert config.pipeline.experimental == {}, ("Dreamverse kwargs leaked into pipeline.experimental: " + f"{config.pipeline.experimental}") def test_refine_enabled_reaches_preset_overrides(self): - config = legacy_from_pretrained_to_config( - "/models/ltx2", _dreamverse_load_kwargs()) + config = legacy_from_pretrained_to_config("/models/ltx2", _dreamverse_load_kwargs()) refine = config.pipeline.preset_overrides.get("refine") or {} assert refine.get("enabled") is True assert refine.get("add_noise") is True @@ -90,17 +87,14 @@ class TestDreamverseLoadKwargsShape: assert refine.get("guidance_scale") == 1.0 def test_refine_assets_reach_component_config(self): - config = legacy_from_pretrained_to_config( - "/models/ltx2", _dreamverse_load_kwargs()) + config = legacy_from_pretrained_to_config("/models/ltx2", _dreamverse_load_kwargs()) assert isinstance(config.pipeline.components, ComponentConfig) - assert (config.pipeline.components.upsampler_weights == - "/models/ltx2-refine") + assert (config.pipeline.components.upsampler_weights == "/models/ltx2-refine") assert config.pipeline.components.lora_path == "/models/ltx2-refine-lora" assert config.pipeline.components.config_root == "/models/ltx2-config" def test_torch_compile_kwargs_reach_typed_fields(self): - config = legacy_from_pretrained_to_config( - "/models/ltx2", _dreamverse_load_kwargs()) + config = legacy_from_pretrained_to_config("/models/ltx2", _dreamverse_load_kwargs()) assert isinstance(config.engine.compile, CompileConfig) assert config.engine.compile.enabled is True assert config.engine.compile.backend == "inductor" @@ -113,21 +107,23 @@ class TestDreamverseLoadKwargsShape: kwargs = _dreamverse_load_kwargs() kwargs["torch_compile_kwargs"] = { **kwargs["torch_compile_kwargs"], - "options": {"epilogue_fusion": True}, + "options": { + "epilogue_fusion": True + }, } config = legacy_from_pretrained_to_config("/models/ltx2", kwargs) assert config.engine.compile.extras == { - "options": {"epilogue_fusion": True}, + "options": { + "epilogue_fusion": True + }, } def test_vae_tiling_reaches_pipeline_selection(self): - config = legacy_from_pretrained_to_config( - "/models/ltx2", _dreamverse_load_kwargs()) + config = legacy_from_pretrained_to_config("/models/ltx2", _dreamverse_load_kwargs()) assert config.pipeline.vae_tiling is True def test_offload_fields_reach_typed_offload_config(self): - config = legacy_from_pretrained_to_config( - "/models/ltx2", _dreamverse_load_kwargs()) + config = legacy_from_pretrained_to_config("/models/ltx2", _dreamverse_load_kwargs()) assert config.engine.offload.dit is False assert config.engine.offload.vae is False assert config.engine.offload.text_encoder is False @@ -188,9 +184,7 @@ class TestDreamverseRequestShape: translation path so Dreamverse callers can opt in.""" request = GenerationRequest( prompt="x", - output=__import__( - "fastvideo.api", fromlist=["OutputConfig"]).OutputConfig( - return_state=True), + output=__import__("fastvideo.api", fromlist=["OutputConfig"]).OutputConfig(return_state=True), ) normalized = normalize_generation_request(request) assert normalized.output.return_state is True diff --git a/fastvideo/tests/contract/test_dynamo_shape.py b/fastvideo/tests/contract/test_dynamo_shape.py index 9432449cb..741f21f53 100644 --- a/fastvideo/tests/contract/test_dynamo_shape.py +++ b/fastvideo/tests/contract/test_dynamo_shape.py @@ -38,7 +38,6 @@ from fastvideo.api import ( SamplingConfig, ) - # ---------------------------------------------------------------------- # Fake Dynamo-side types (the real ones live in dynamo.llm and dynamo.*) # We mimic their dict shape so the adapter code under test is realistic. @@ -87,8 +86,7 @@ def _parse_size(size: str | None) -> tuple[int, int]: return int(w), int(h) -def nv_request_to_generation_request( - request: dict[str, Any], ) -> GenerationRequest: +def nv_request_to_generation_request(request: dict[str, Any], ) -> GenerationRequest: """Translate Dynamo's request shape into FastVideo's typed request. This function is the template integrators copy into the Dynamo repo. @@ -150,7 +148,9 @@ class _MockFastVideoHandler: await asyncio.sleep(0) # model call stub elapsed = time.perf_counter() - t0 yield { - "data": [{"b64_json": self._fake_result["b64_json"]}], + "data": [{ + "b64_json": self._fake_result["b64_json"] + }], "inference_time_s": elapsed, "model": request.get("model"), "nvext": { @@ -184,13 +184,11 @@ class TestDynamoAdapterShape: assert req.negative_prompt == "blurry" def test_nvext_num_frames_overrides_seconds_times_fps(self): - req = nv_request_to_generation_request( - _nv_create_video_request(nvext={"num_frames": 33})) + req = nv_request_to_generation_request(_nv_create_video_request(nvext={"num_frames": 33})) assert req.sampling.num_frames == 33 def test_input_reference_maps_to_image_path(self): - req = nv_request_to_generation_request( - _nv_create_video_request(input_reference="/tmp/init.png")) + req = nv_request_to_generation_request(_nv_create_video_request(input_reference="/tmp/init.png")) assert req.inputs.image_path == "/tmp/init.png" def test_continuation_state_round_trips_through_adapter(self): @@ -201,8 +199,7 @@ class TestDynamoAdapterShape: "segment_index": 2, }, } - req = nv_request_to_generation_request( - _nv_create_video_request(nvext={"continuation_state": state})) + req = nv_request_to_generation_request(_nv_create_video_request(nvext={"continuation_state": state})) assert isinstance(req.state, ContinuationState) assert req.state.kind == "ltx2.v1" assert req.state.payload["segment_index"] == 2 @@ -215,8 +212,7 @@ class TestDynamoHandlerContract: async def run(): out = [] - async for chunk in handler.generate( - _nv_create_video_request(), _FakeContext()): + async for chunk in handler.generate(_nv_create_video_request(), _FakeContext()): out.append(chunk) return out @@ -235,22 +231,21 @@ class TestDynamoHandlerContract: handler = _MockFastVideoHandler({"b64_json": "abc"}) state_in = { "kind": "ltx2.v1", - "payload": {"schema_version": 1, "segment_index": 4}, + "payload": { + "schema_version": 1, + "segment_index": 4 + }, } async def run(): - async for chunk in handler.generate( - _nv_create_video_request( - nvext={"continuation_state": state_in}), - _FakeContext()): + async for chunk in handler.generate(_nv_create_video_request(nvext={"continuation_state": state_in}), + _FakeContext()): return chunk raise AssertionError("no chunk emitted") chunk = asyncio.run(run()) assert chunk["nvext"]["continuation_state"]["kind"] == "ltx2.v1" - assert ( - chunk["nvext"]["continuation_state"]["payload"]["segment_index"] - == 4) + assert (chunk["nvext"]["continuation_state"]["payload"]["segment_index"] == 4) class TestNoInternalImports: @@ -284,20 +279,15 @@ class TestNoInternalImports: adapter_source = inspect.getsource(nv_request_to_generation_request) for banned in self._BANNED_PREFIXES: - assert banned not in adapter_source, ( - f"Dynamo adapter template must not depend on {banned}*; " - "found a reference in the adapter function body.") + assert banned not in adapter_source, (f"Dynamo adapter template must not depend on {banned}*; " + "found a reference in the adapter function body.") def test_public_types_are_stable_imports(self): # These lines match the exact import snippet in docs/design/ # server_contracts/dynamo.md; any rename breaks CI. from fastvideo import VideoGenerator # noqa: F401 from fastvideo.api import ( # noqa: F401 - ContinuationState, - GenerationRequest, - InputConfig, - OutputConfig, - SamplingConfig, + ContinuationState, GenerationRequest, InputConfig, OutputConfig, SamplingConfig, ) @pytest.mark.skipif( @@ -315,8 +305,7 @@ class TestNoInternalImports: "SamplingConfig", } missing = required - set(api.__all__) - assert not missing, ( - f"fastvideo.api.__all__ missing Dynamo-contract symbols: {missing}") + assert not missing, (f"fastvideo.api.__all__ missing Dynamo-contract symbols: {missing}") def _read_own_source() -> str: diff --git a/fastvideo/tests/contract/test_generate_async.py b/fastvideo/tests/contract/test_generate_async.py index 3d2c35e7f..50aded49d 100644 --- a/fastvideo/tests/contract/test_generate_async.py +++ b/fastvideo/tests/contract/test_generate_async.py @@ -46,7 +46,8 @@ class _FakeVideoGenerator: self.executor = _FakeExecutor() def _generate_request_impl( - self, request: GenerationRequest, + self, + request: GenerationRequest, ) -> GenerationResult | list[GenerationResult]: return self._result @@ -59,18 +60,15 @@ class _FakeVideoGenerator: from fastvideo.entrypoints.video_generator import VideoGenerator instance = cls(result) - instance.generate_async = VideoGenerator.generate_async.__get__( - instance, cls) - instance.default_health_check_request = ( - VideoGenerator.default_health_check_request) + instance.generate_async = VideoGenerator.generate_async.__get__(instance, cls) + instance.default_health_check_request = (VideoGenerator.default_health_check_request) return instance def _make_request(**overrides: Any) -> GenerationRequest: base = GenerationRequest( prompt="hi", - sampling=SamplingConfig(num_inference_steps=4, height=64, width=64, - num_frames=8), + sampling=SamplingConfig(num_inference_steps=4, height=64, width=64, num_frames=8), output=OutputConfig(save_video=False, return_frames=False), ) for key, value in overrides.items(): @@ -152,8 +150,7 @@ class TestFinalEventShape: assert final.metadata["generation_time"] == 0.01 def test_carries_continuation_state_when_present(self): - state = ContinuationState( - kind="ltx2.v1", payload={"schema_version": 1, "segment_index": 3}) + state = ContinuationState(kind="ltx2.v1", payload={"schema_version": 1, "segment_index": 3}) result = _make_result(state=state) gen = _FakeVideoGenerator.bind(result) @@ -235,12 +232,7 @@ class TestDynamoStyleHandlerIntegration: def test_async_handler_yields_typed_events_without_internal_imports(self): # The import set is exactly what the Dynamo backend uses. from fastvideo.api import ( # noqa: F401 - ContinuationState, - GenerationRequest, - InputConfig, - OutputConfig, - SamplingConfig, - VideoFinalEvent, + ContinuationState, GenerationRequest, InputConfig, OutputConfig, SamplingConfig, VideoFinalEvent, VideoProgressEvent, ) @@ -250,9 +242,7 @@ class TestDynamoStyleHandlerIntegration: # Adapter: dict -> typed request. req = GenerationRequest( prompt=request_dict["prompt"], - sampling=SamplingConfig( - num_inference_steps=1, num_frames=8, - height=256, width=256), + sampling=SamplingConfig(num_inference_steps=1, num_frames=8, height=256, width=256), output=OutputConfig(save_video=False, return_frames=False), ) async for event in gen.generate_async(req): diff --git a/fastvideo/tests/contract/test_ltx2_compile_conditions.py b/fastvideo/tests/contract/test_ltx2_compile_conditions.py index 89d68cbd6..179986fb9 100644 --- a/fastvideo/tests/contract/test_ltx2_compile_conditions.py +++ b/fastvideo/tests/contract/test_ltx2_compile_conditions.py @@ -19,9 +19,8 @@ from fastvideo.models.dits import ltx2 def test_ltx2_model_exposes_compile_conditions(): conditions = ltx2.LTX2Transformer3DModel._compile_conditions - assert conditions, ( - "LTX2Transformer3DModel._compile_conditions is empty: " - "torch.compile falls back to a whole-model fullgraph") + assert conditions, ("LTX2Transformer3DModel._compile_conditions is empty: " + "torch.compile falls back to a whole-model fullgraph") assert is_ltx2_blocks in conditions assert LTX2VideoConfig()._compile_conditions == conditions diff --git a/fastvideo/tests/contract/test_modal_fa4_policy.py b/fastvideo/tests/contract/test_modal_fa4_policy.py index e5d878393..1d0fa0e12 100644 --- a/fastvideo/tests/contract/test_modal_fa4_policy.py +++ b/fastvideo/tests/contract/test_modal_fa4_policy.py @@ -21,12 +21,8 @@ def _function_strings(path: Path, function_name: str) -> str: tree = ast.parse(source) for node in tree.body: if isinstance(node, ast.FunctionDef) and node.name == function_name: - return "\n".join( - child.value - for child in ast.walk(node) - if isinstance(child, ast.Constant) - and isinstance(child.value, str) - ) + return "\n".join(child.value for child in ast.walk(node) + if isinstance(child, ast.Constant) and isinstance(child.value, str)) raise AssertionError(f"{function_name} not found in {path}") @@ -52,8 +48,7 @@ def test_performance_identity_env_reaches_modal_runtime(): def test_performance_lane_classifies_pull_requests_before_main(): function_strings = _function_strings(PR_TEST, "run_performance_tests") - assert function_strings.index( - "BUILDKITE_PULL_REQUEST") < function_strings.index("BUILDKITE_BRANCH") + assert function_strings.index("BUILDKITE_PULL_REQUEST") < function_strings.index("BUILDKITE_BRANCH") def test_vsa_training_lane_uses_strict_h100_pair(): diff --git a/fastvideo/tests/dataset/test_load_image_retry.py b/fastvideo/tests/dataset/test_load_image_retry.py index 7d7414fb3..b3513d575 100644 --- a/fastvideo/tests/dataset/test_load_image_retry.py +++ b/fastvideo/tests/dataset/test_load_image_retry.py @@ -35,8 +35,7 @@ def test_load_image_retries_transient_error_then_succeeds(monkeypatch): monkeypatch.setattr(vision_utils.time, "sleep", lambda _s: None) get = mock.Mock(side_effect=[ # The shape that broke build 4202: connection dropped mid-body. - requests.exceptions.ChunkedEncodingError( - "IncompleteRead(1024 bytes read, 2048 more expected)"), + requests.exceptions.ChunkedEncodingError("IncompleteRead(1024 bytes read, 2048 more expected)"), _good_response(), ]) monkeypatch.setattr(vision_utils.requests, "get", get) @@ -53,8 +52,7 @@ def test_load_image_retries_transient_error_then_succeeds(monkeypatch): def test_load_image_raises_after_exhausting_retries(monkeypatch): sleeps: list[float] = [] monkeypatch.setattr(vision_utils.time, "sleep", sleeps.append) - get = mock.Mock( - side_effect=requests.exceptions.ConnectionError("connection reset")) + get = mock.Mock(side_effect=requests.exceptions.ConnectionError("connection reset")) monkeypatch.setattr(vision_utils.requests, "get", get) with pytest.raises(requests.exceptions.ConnectionError): diff --git a/fastvideo/tests/dataset/test_parquet_dataset_map_style.py b/fastvideo/tests/dataset/test_parquet_dataset_map_style.py index 04dd55766..e899cd0ae 100644 --- a/fastvideo/tests/dataset/test_parquet_dataset_map_style.py +++ b/fastvideo/tests/dataset/test_parquet_dataset_map_style.py @@ -40,7 +40,11 @@ def test_parse_data_path_specs_accepts_path_list() -> None: def test_parse_data_path_specs_drops_non_positive_repeats() -> None: - assert _parse_data_path_specs({"data/a": 0, "data/b": -1, "data/c": 2}) == [ + assert _parse_data_path_specs({ + "data/a": 0, + "data/b": -1, + "data/c": 2 + }) == [ ("data/c", 2), ] assert _parse_data_path_specs("data/a:0,data/b:2") == [("data/b", 2)] @@ -69,7 +73,7 @@ def _write_root_cache(dataset_root, filename: str, length: int) -> str: parquet_file = dataset_root / filename parquet_file.touch() with (cache_dir / "file_info.pkl").open("wb") as f: - pickle.dump(((str(parquet_file),), (length,)), f) + pickle.dump(((str(parquet_file), ), (length, )), f) return str(parquet_file) diff --git a/fastvideo/tests/dataset/test_parquet_io.py b/fastvideo/tests/dataset/test_parquet_io.py index 744f26cdb..28b01a99e 100644 --- a/fastvideo/tests/dataset/test_parquet_io.py +++ b/fastvideo/tests/dataset/test_parquet_io.py @@ -104,5 +104,3 @@ def test_writer_parallel_workers(tmp_path: Path): assert len(files) == 4 total_rows = sum(pq.read_table(str(f)).num_rows for f in files) assert total_rows == 40 - - diff --git a/fastvideo/tests/distributed/test_sp_hunyuanvideo.py b/fastvideo/tests/distributed/test_sp_hunyuanvideo.py index fa598510a..78c95a6c2 100644 --- a/fastvideo/tests/distributed/test_sp_hunyuanvideo.py +++ b/fastvideo/tests/distributed/test_sp_hunyuanvideo.py @@ -91,9 +91,7 @@ def _build_inputs( return hidden_states, encoder_hidden_states, timestep -def _collect_named_parameter_grads( - model: torch.nn.Module, -) -> dict[str, torch.Tensor]: +def _collect_named_parameter_grads(model: torch.nn.Module, ) -> dict[str, torch.Tensor]: named_grads: dict[str, torch.Tensor] = {} for name, param in model.named_parameters(): grad = param.grad @@ -124,9 +122,7 @@ def _assert_finite(name: str, tensor: torch.Tensor) -> None: return nan_count = int(torch.isnan(tensor).sum().item()) inf_count = int(torch.isinf(tensor).sum().item()) - raise RuntimeError( - f"{name} contains non-finite values (nan={nan_count}, inf={inf_count})" - ) + raise RuntimeError(f"{name} contains non-finite values (nan={nan_count}, inf={inf_count})") def _run_worker(mode: str, output_path: Path) -> None: @@ -158,9 +154,9 @@ def _run_worker(mode: str, output_path: Path) -> None: forward_batch = ForwardBatch(data_type="dummy") with set_forward_context( - current_timestep=0, - attn_metadata=None, - forward_batch=forward_batch, + current_timestep=0, + attn_metadata=None, + forward_batch=forward_batch, ): output = model( hidden_states=hidden_states, @@ -214,20 +210,16 @@ def _run_torchrun( env["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA" process = subprocess.run(cmd, capture_output=True, text=True, env=env) if process.returncode != 0: - raise RuntimeError( - f"{mode} worker failed with code {process.returncode}\n" - f"STDOUT:\n{process.stdout}\n" - f"STDERR:\n{process.stderr}" - ) + raise RuntimeError(f"{mode} worker failed with code {process.returncode}\n" + f"STDOUT:\n{process.stdout}\n" + f"STDERR:\n{process.stderr}") def test_sp_gradient_matches_single_rank(tmp_path: Path) -> None: if not torch.cuda.is_available(): pytest.skip("This test requires CUDA.") if torch.cuda.device_count() < SP_WORLD_SIZE: - pytest.skip( - f"This test requires at least {SP_WORLD_SIZE} CUDA devices." - ) + pytest.skip(f"This test requires at least {SP_WORLD_SIZE} CUDA devices.") script_path = Path(__file__).resolve() single_path = tmp_path / "single_rank_grads.pt" @@ -270,17 +262,10 @@ def test_sp_gradient_matches_single_rank(tmp_path: Path) -> None: max_abs_diff = (single_grad - sp_grad).abs().max().item() different_params.append(f"{name} (max_abs_diff={max_abs_diff:.3e})") - report = ( - f"single-only params ({len(single_only)}):\n" - + "\n".join(single_only) - + f"\n\nsp-only params ({len(sp_only)}):\n" - + "\n".join(sp_only) - + "\n\n" - f"same params ({len(same_params)}):\n" - + "\n".join(same_params) - + f"\n\ndifferent params ({len(different_params)}):\n" - + "\n".join(different_params) - ) + report = (f"single-only params ({len(single_only)}):\n" + "\n".join(single_only) + + f"\n\nsp-only params ({len(sp_only)}):\n" + "\n".join(sp_only) + "\n\n" + f"same params ({len(same_params)}):\n" + "\n".join(same_params) + + f"\n\ndifferent params ({len(different_params)}):\n" + "\n".join(different_params)) print(report) assert not single_only and not sp_only and not different_params, report diff --git a/fastvideo/tests/distributed/test_sp_ltx2.py b/fastvideo/tests/distributed/test_sp_ltx2.py index afbdfbb5e..8ec8f7a70 100644 --- a/fastvideo/tests/distributed/test_sp_ltx2.py +++ b/fastvideo/tests/distributed/test_sp_ltx2.py @@ -98,11 +98,8 @@ def _build_inputs( ).to(device) patch_t, patch_h, patch_w = config.patch_size - sequence_len = ( - (hidden_states.shape[2] // patch_t) - * (hidden_states.shape[3] // patch_h) - * (hidden_states.shape[4] // patch_w) - ) + sequence_len = ((hidden_states.shape[2] // patch_t) * (hidden_states.shape[3] // patch_h) * + (hidden_states.shape[4] // patch_w)) timestep = torch.full( (1, sequence_len), 10, @@ -112,9 +109,7 @@ def _build_inputs( return hidden_states, encoder_hidden_states, timestep -def _collect_named_parameter_grads( - model: torch.nn.Module, -) -> dict[str, torch.Tensor]: +def _collect_named_parameter_grads(model: torch.nn.Module, ) -> dict[str, torch.Tensor]: named_grads: dict[str, torch.Tensor] = {} for name, param in model.named_parameters(): grad = param.grad @@ -145,9 +140,7 @@ def _assert_finite(name: str, tensor: torch.Tensor) -> None: return nan_count = int(torch.isnan(tensor).sum().item()) inf_count = int(torch.isinf(tensor).sum().item()) - raise RuntimeError( - f"{name} contains non-finite values (nan={nan_count}, inf={inf_count})" - ) + raise RuntimeError(f"{name} contains non-finite values (nan={nan_count}, inf={inf_count})") def _run_worker(mode: str, output_path: Path) -> None: @@ -179,9 +172,9 @@ def _run_worker(mode: str, output_path: Path) -> None: forward_batch = ForwardBatch(data_type="dummy") with set_forward_context( - current_timestep=0, - attn_metadata=None, - forward_batch=forward_batch, + current_timestep=0, + attn_metadata=None, + forward_batch=forward_batch, ): output = model( hidden_states=hidden_states, @@ -235,20 +228,16 @@ def _run_torchrun( env["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA" process = subprocess.run(cmd, capture_output=True, text=True, env=env) if process.returncode != 0: - raise RuntimeError( - f"{mode} worker failed with code {process.returncode}\n" - f"STDOUT:\n{process.stdout}\n" - f"STDERR:\n{process.stderr}" - ) + raise RuntimeError(f"{mode} worker failed with code {process.returncode}\n" + f"STDOUT:\n{process.stdout}\n" + f"STDERR:\n{process.stderr}") def test_sp_gradient_matches_single_rank(tmp_path: Path) -> None: if not torch.cuda.is_available(): pytest.skip("This test requires CUDA.") if torch.cuda.device_count() < SP_WORLD_SIZE: - pytest.skip( - f"This test requires at least {SP_WORLD_SIZE} CUDA devices." - ) + pytest.skip(f"This test requires at least {SP_WORLD_SIZE} CUDA devices.") script_path = Path(__file__).resolve() single_path = tmp_path / "single_rank_grads.pt" @@ -291,17 +280,10 @@ def test_sp_gradient_matches_single_rank(tmp_path: Path) -> None: max_abs_diff = (single_grad - sp_grad).abs().max().item() different_params.append(f"{name} (max_abs_diff={max_abs_diff:.3e})") - report = ( - f"single-only params ({len(single_only)}):\n" - + "\n".join(single_only) - + f"\n\nsp-only params ({len(sp_only)}):\n" - + "\n".join(sp_only) - + "\n\n" - f"same params ({len(same_params)}):\n" - + "\n".join(same_params) - + f"\n\ndifferent params ({len(different_params)}):\n" - + "\n".join(different_params) - ) + report = (f"single-only params ({len(single_only)}):\n" + "\n".join(single_only) + + f"\n\nsp-only params ({len(sp_only)}):\n" + "\n".join(sp_only) + "\n\n" + f"same params ({len(same_params)}):\n" + "\n".join(same_params) + + f"\n\ndifferent params ({len(different_params)}):\n" + "\n".join(different_params)) print(report) assert not single_only and not sp_only and not different_params, report diff --git a/fastvideo/tests/distributed/test_sp_wan.py b/fastvideo/tests/distributed/test_sp_wan.py index 96de18d4b..83369131c 100644 --- a/fastvideo/tests/distributed/test_sp_wan.py +++ b/fastvideo/tests/distributed/test_sp_wan.py @@ -100,9 +100,7 @@ def _build_inputs( return hidden_states, encoder_hidden_states, timestep -def _collect_named_parameter_grads( - model: torch.nn.Module, -) -> dict[str, torch.Tensor]: +def _collect_named_parameter_grads(model: torch.nn.Module, ) -> dict[str, torch.Tensor]: named_grads: dict[str, torch.Tensor] = {} for name, param in model.named_parameters(): grad = param.grad @@ -133,9 +131,7 @@ def _assert_finite(name: str, tensor: torch.Tensor) -> None: return nan_count = int(torch.isnan(tensor).sum().item()) inf_count = int(torch.isinf(tensor).sum().item()) - raise RuntimeError( - f"{name} contains non-finite values (nan={nan_count}, inf={inf_count})" - ) + raise RuntimeError(f"{name} contains non-finite values (nan={nan_count}, inf={inf_count})") def _run_worker(mode: str, output_path: Path) -> None: @@ -168,9 +164,9 @@ def _run_worker(mode: str, output_path: Path) -> None: forward_batch = ForwardBatch(data_type="dummy") with set_forward_context( - current_timestep=0, - attn_metadata=None, - forward_batch=forward_batch, + current_timestep=0, + attn_metadata=None, + forward_batch=forward_batch, ): output = model( hidden_states=hidden_states, @@ -224,20 +220,16 @@ def _run_torchrun( env["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA" process = subprocess.run(cmd, capture_output=True, text=True, env=env) if process.returncode != 0: - raise RuntimeError( - f"{mode} worker failed with code {process.returncode}\n" - f"STDOUT:\n{process.stdout}\n" - f"STDERR:\n{process.stderr}" - ) + raise RuntimeError(f"{mode} worker failed with code {process.returncode}\n" + f"STDOUT:\n{process.stdout}\n" + f"STDERR:\n{process.stderr}") def test_sp_gradient_matches_single_rank(tmp_path: Path) -> None: if not torch.cuda.is_available(): pytest.skip("This test requires CUDA.") if torch.cuda.device_count() < SP_WORLD_SIZE: - pytest.skip( - f"This test requires at least {SP_WORLD_SIZE} CUDA devices." - ) + pytest.skip(f"This test requires at least {SP_WORLD_SIZE} CUDA devices.") script_path = Path(__file__).resolve() single_path = tmp_path / "single_rank_grads.pt" @@ -280,17 +272,10 @@ def test_sp_gradient_matches_single_rank(tmp_path: Path) -> None: max_abs_diff = (single_grad - sp_grad).abs().max().item() different_params.append(f"{name} (max_abs_diff={max_abs_diff:.3e})") - report = ( - f"single-only params ({len(single_only)}):\n" - + "\n".join(single_only) - + f"\n\nsp-only params ({len(sp_only)}):\n" - + "\n".join(sp_only) - + "\n\n" - f"same params ({len(same_params)}):\n" - + "\n".join(same_params) - + f"\n\ndifferent params ({len(different_params)}):\n" - + "\n".join(different_params) - ) + report = (f"single-only params ({len(single_only)}):\n" + "\n".join(single_only) + + f"\n\nsp-only params ({len(sp_only)}):\n" + "\n".join(sp_only) + "\n\n" + f"same params ({len(same_params)}):\n" + "\n".join(same_params) + + f"\n\ndifferent params ({len(different_params)}):\n" + "\n".join(different_params)) print(report) assert not single_only and not sp_only and not different_params, report @@ -302,7 +287,8 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--output", type=str, default=None) return parser.parse_args() -# pytest -sv fastvideo/tests/distributed/test_sp_wan.py + +# pytest -sv fastvideo/tests/distributed/test_sp_wan.py if __name__ == "__main__": args = _parse_args() if not args.sp_grad_worker: diff --git a/fastvideo/tests/encoders/test_clip_encoder.py b/fastvideo/tests/encoders/test_clip_encoder.py index 266a9ef5d..273667db2 100644 --- a/fastvideo/tests/encoders/test_clip_encoder.py +++ b/fastvideo/tests/encoders/test_clip_encoder.py @@ -97,9 +97,8 @@ def test_clip_encoder(): logger.info("Testing prompt: '%s'", prompt) # Tokenize the prompt - tokens = tokenizer(prompt, padding="max_length", max_length=77, truncation=True, return_tensors="pt").to( - device - ) + tokens = tokenizer(prompt, padding="max_length", max_length=77, truncation=True, + return_tensors="pt").to(device) # Get embeddings from our implementation outputs1 = model1(input_ids=tokens.input_ids, output_hidden_states=True) @@ -120,15 +119,13 @@ def test_clip_encoder(): # print("last_hidden_state2", last_hidden_state2) assert last_hidden_state1.shape == last_hidden_state2.shape, ( - f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}" - ) + f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}") # Compare pooler outputs pooler_output1 = outputs1.pooler_output pooler_output2 = outputs2.pooler_output assert pooler_output1.shape == pooler_output2.shape, ( - f"Pooler output shapes don't match: {pooler_output1.shape} vs {pooler_output2.shape}" - ) + f"Pooler output shapes don't match: {pooler_output1.shape} vs {pooler_output2.shape}") assert_close(pooler_output1, pooler_output2, atol=1e-2, rtol=1e-3) assert_close(last_hidden_state1, last_hidden_state2, atol=1e-2, rtol=1e-3) diff --git a/fastvideo/tests/encoders/test_gemma_connector_register_padding.py b/fastvideo/tests/encoders/test_gemma_connector_register_padding.py index fd21cab5a..a726e2841 100644 --- a/fastvideo/tests/encoders/test_gemma_connector_register_padding.py +++ b/fastvideo/tests/encoders/test_gemma_connector_register_padding.py @@ -24,8 +24,8 @@ class _Stub: # Register i holds -(i+1) in every channel: distinct from any hidden # state below, and distinct from each other, so an assertion can pin # not just "a register landed here" but *which* register landed here. - self.learnable_registers = -(torch.arange(num_registers, dtype=torch.float32) + - 1.0).unsqueeze(-1).expand(num_registers, hidden).contiguous() + self.learnable_registers = -(torch.arange(num_registers, dtype=torch.float32) + 1.0).unsqueeze(-1).expand( + num_registers, hidden).contiguous() def _run(hidden_states: torch.Tensor, valid_lengths: list[int], num_registers: int | None = None): @@ -47,10 +47,8 @@ def test_ragged_batch_left_aligns_each_row_independently() -> None: """ seq_len, hidden = 4, 2 # Token values encode (row, position) so misplacement is visible. - hidden_states = torch.tensor( - [[[10.0, 10.0], [11.0, 11.0], [12.0, 12.0], [13.0, 13.0]], - [[20.0, 20.0], [21.0, 21.0], [22.0, 22.0], [23.0, 23.0]]], - ) + hidden_states = torch.tensor([[[10.0, 10.0], [11.0, 11.0], [12.0, 12.0], [13.0, 13.0]], + [[20.0, 20.0], [21.0, 21.0], [22.0, 22.0], [23.0, 23.0]]], ) out, _ = _run(hidden_states, valid_lengths=[3, 1]) # Row 0's valid tokens are the last 3 (11, 12, 13), left-aligned in order. @@ -91,10 +89,8 @@ def test_registers_tile_across_the_sequence() -> None: def test_single_row_matches_ragged_batch_row() -> None: """A row's result must not depend on what else is in the batch.""" - hidden_states = torch.tensor( - [[[10.0, 10.0], [11.0, 11.0], [12.0, 12.0], [13.0, 13.0]], - [[20.0, 20.0], [21.0, 21.0], [22.0, 22.0], [23.0, 23.0]]], - ) + hidden_states = torch.tensor([[[10.0, 10.0], [11.0, 11.0], [12.0, 12.0], [13.0, 13.0]], + [[20.0, 20.0], [21.0, 21.0], [22.0, 22.0], [23.0, 23.0]]], ) batched, _ = _run(hidden_states, valid_lengths=[3, 1]) alone, _ = _run(hidden_states[1:2], valid_lengths=[1]) assert torch.equal(batched[1], alone[0]) diff --git a/fastvideo/tests/encoders/test_hyt5_encoder.py b/fastvideo/tests/encoders/test_hyt5_encoder.py index f072d7444..e1d9e4f8c 100644 --- a/fastvideo/tests/encoders/test_hyt5_encoder.py +++ b/fastvideo/tests/encoders/test_hyt5_encoder.py @@ -41,14 +41,10 @@ def test_t5_encoder(t5_model_paths_and_config): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") precision_str = "fp32" precision = PRECISION_TO_TYPE[precision_str] - model1 = T5EncoderModel.from_pretrained(text_encoder_path).to( - precision).to(device).eval() + model1 = T5EncoderModel.from_pretrained(text_encoder_path).to(precision).to(device).eval() tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - - args = FastVideoArgs(model_path=text_encoder_path, - pipeline_config=pipeline_config, - pin_cpu_memory=False) + args = FastVideoArgs(model_path=text_encoder_path, pipeline_config=pipeline_config, pin_cpu_memory=False) loader = TextEncoderLoader() model2 = loader.load(text_encoder_path, args) model2 = model2.to(precision) @@ -68,7 +64,7 @@ def test_t5_encoder(t5_model_paths_and_config): "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_0.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_1.weight",\ "encoder.block.{}.layer.1.DenseReluDense.wo.weight", \ "encoder.block.{}.layer.1.layer_norm.weight", "encoder.final_layer_norm.weight"] - + for idx in range(hf_config.num_hidden_layers): for w in weights: name1 = w.format(idx) @@ -77,13 +73,9 @@ def test_t5_encoder(t5_model_paths_and_config): p2 = params2[name2] p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1) assert_close(p1, p2, atol=1e-4, rtol=1e-4) - # Test with some sample prompts - prompts = [ - "Once upon a time", "The quick brown fox jumps over", - "In a galaxy far, far away" - ] + prompts = ["Once upon a time", "The quick brown fox jumps over", "In a galaxy far, far away"] logger.info("Testing T5 encoder with sample prompts") @@ -103,8 +95,7 @@ def test_t5_encoder(t5_model_paths_and_config): # filter out padding input_ids # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1] # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1] - outputs1 = model1(input_ids=tokens.input_ids, - attention_mask=tokens.attention_mask.float())[0] + outputs1 = model1(input_ids=tokens.input_ids, attention_mask=tokens.attention_mask.float())[0] print("--------------------------------") logger.info("Testing model2") @@ -122,15 +113,11 @@ def test_t5_encoder(t5_model_paths_and_config): assert last_hidden_state1.shape == last_hidden_state2.shape, \ f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}" - max_diff_hidden = torch.max( - torch.abs(last_hidden_state1 - last_hidden_state2)) - mean_diff_hidden = torch.mean( - torch.abs(last_hidden_state1 - last_hidden_state2)) + max_diff_hidden = torch.max(torch.abs(last_hidden_state1 - last_hidden_state2)) + mean_diff_hidden = torch.mean(torch.abs(last_hidden_state1 - last_hidden_state2)) - logger.info("Maximum difference in last hidden states: %s", - max_diff_hidden.item()) - logger.info("Mean difference in last hidden states: %s", - mean_diff_hidden.item()) + logger.info("Maximum difference in last hidden states: %s", max_diff_hidden.item()) + logger.info("Mean difference in last hidden states: %s", mean_diff_hidden.item()) logger.info("Max memory allocated: %s GB", torch.cuda.max_memory_allocated() / 1024**3) # Check if outputs are similar (allowing for small numerical differences) assert mean_diff_hidden < 1e-4, \ diff --git a/fastvideo/tests/encoders/test_llama_encoder.py b/fastvideo/tests/encoders/test_llama_encoder.py index e04ac5ba4..728405fee 100644 --- a/fastvideo/tests/encoders/test_llama_encoder.py +++ b/fastvideo/tests/encoders/test_llama_encoder.py @@ -71,8 +71,7 @@ def test_llama_encoder(): assert torch.allclose( model1.embed_tokens.weight, model2.embed_tokens.weight.to_local().to(device) - if isinstance(model2.embed_tokens.weight, DTensor) - else model2.embed_tokens.weight.to(device), + if isinstance(model2.embed_tokens.weight, DTensor) else model2.embed_tokens.weight.to(device), ) weights = ["layers.{}.input_layernorm.weight", "layers.{}.post_attention_layernorm.weight"] for layer_idx in range(hf_config.num_hidden_layers): @@ -119,9 +118,8 @@ def test_llama_encoder(): logger.info("Testing prompt: '%s'", prompt) # Tokenize the prompt - tokens = tokenizer(prompt, padding="max_length", max_length=512, truncation=True, return_tensors="pt").to( - device - ) + tokens = tokenizer(prompt, padding="max_length", max_length=512, truncation=True, + return_tensors="pt").to(device) # Get outputs from our implementation # filter out padding input_ids @@ -133,15 +131,14 @@ def test_llama_encoder(): # Get outputs from HuggingFace implementation with set_forward_context(current_timestep=0, attn_metadata=None): - outputs2 = model2( - input_ids=tokens.input_ids, attention_mask=tokens.attention_mask, output_hidden_states=True - ) + outputs2 = model2(input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True) # Compare last hidden states last_hidden_state1 = outputs1.last_hidden_state[tokens.attention_mask == 1] last_hidden_state2 = outputs2.last_hidden_state[tokens.attention_mask == 1] assert last_hidden_state1.shape == last_hidden_state2.shape, ( - f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}" - ) + f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}") assert_close(last_hidden_state1, last_hidden_state2, atol=1e-1, rtol=1e-4) diff --git a/fastvideo/tests/encoders/test_qwen2_5_encoder.py b/fastvideo/tests/encoders/test_qwen2_5_encoder.py index 1386359e4..85d3de114 100644 --- a/fastvideo/tests/encoders/test_qwen2_5_encoder.py +++ b/fastvideo/tests/encoders/test_qwen2_5_encoder.py @@ -19,6 +19,7 @@ logger = init_logger(__name__) os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29505" + @pytest.fixture def qwen_model_path_and_config(): base_model_path = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" @@ -27,6 +28,7 @@ def qwen_model_path_and_config(): tokenizer_path = os.path.join(model_path, "tokenizer") return text_encoder_path, tokenizer_path, Hunyuan15T2V480PConfig() + @pytest.mark.usefixtures("distributed_setup") def test_qwen2_5_encoder(qwen_model_path_and_config): text_encoder_path, tokenizer_path, pipeline_config = qwen_model_path_and_config @@ -37,19 +39,16 @@ def test_qwen2_5_encoder(qwen_model_path_and_config): # Qwen2.5-VL default dtype is usually bf16 precision_str = "fp32" precision = PRECISION_TO_TYPE[precision_str] - + logger.info(f"Using precision: {precision_str}") # Load HF model (Base model) - model1 = Qwen2_5_VLTextModel.from_pretrained(text_encoder_path).to( - precision).to(device).eval() + model1 = Qwen2_5_VLTextModel.from_pretrained(text_encoder_path).to(precision).to(device).eval() tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) # Load FastVideo model - args = FastVideoArgs(model_path=text_encoder_path, - pipeline_config=pipeline_config, - pin_cpu_memory=False) - + args = FastVideoArgs(model_path=text_encoder_path, pipeline_config=pipeline_config, pin_cpu_memory=False) + loader = TextEncoderLoader() model2 = loader.load(text_encoder_path, args) model2 = model2.to(precision) @@ -59,7 +58,7 @@ def test_qwen2_5_encoder(qwen_model_path_and_config): logger.info("Comparing model weights for sanity check...") params1 = dict(model1.named_parameters()) params2 = dict(model2.named_parameters()) - + logger.info("Model1 has %s parameters", len(params1)) logger.info("Model2 has %s parameters", len(params2)) @@ -67,15 +66,12 @@ def test_qwen2_5_encoder(qwen_model_path_and_config): # or simple linear layers if names match. # Note: FastVideo uses QKVParallelLinear, so q_proj, k_proj, v_proj are merged. # HF Qwen2_5_VL uses separate projections? No, usually they are separate nn.Linear in HF. - + weights_to_check = [ - "norm.weight", - "layers.{}.self_attn.o_proj.weight", - "layers.{}.input_layernorm.weight", - "layers.{}.post_attention_layernorm.weight", - "layers.{}.mlp.down_proj.weight" + "norm.weight", "layers.{}.self_attn.o_proj.weight", "layers.{}.input_layernorm.weight", + "layers.{}.post_attention_layernorm.weight", "layers.{}.mlp.down_proj.weight" ] - + for idx in range(hf_config.num_hidden_layers): for w in weights_to_check: name1 = w.format(idx) @@ -83,64 +79,57 @@ def test_qwen2_5_encoder(qwen_model_path_and_config): p1 = params1[name1] p2 = params2[name2] p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1) - + # Check shape assert p1.shape == p2.shape, f"Shape mismatch for {w}: {p1.shape} vs {p2.shape}" - + # Check values assert_close(p1, p2, atol=1e-7, rtol=1e-7, msg=f"Weight mismatch for {w}") # Test with sample prompts - prompts = [ - "Hello world", - "The quick brown fox jumps over the lazy dog." - ] - + prompts = ["Hello world", "The quick brown fox jumps over the lazy dog."] + logger.info("Testing with sample prompts") - + with torch.no_grad(): for prompt in prompts: logger.info(f"Prompt: {prompt}") - tokens = tokenizer(prompt, return_tensors="pt", padding="max_length", max_length=1000, truncation=True).to(device) - + tokens = tokenizer(prompt, return_tensors="pt", padding="max_length", max_length=1000, + truncation=True).to(device) + # HF Forward # AutoModel for Qwen2.5-VL usually returns BaseModelOutputWithPast - outputs1 = model1( - input_ids=tokens.input_ids, - attention_mask=tokens.attention_mask, - output_hidden_states=True - ).hidden_states[-3] - + outputs1 = model1(input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True).hidden_states[-3] + # FastVideo Forward with set_forward_context(current_timestep=0, attn_metadata=None): - outputs2 = model2( - input_ids=tokens.input_ids, - attention_mask=tokens.attention_mask, - output_hidden_states=True - ).hidden_states[-3] - + outputs2 = model2(input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True).hidden_states[-3] + # Compare # Filter padding for comparison if needed, but here we just check raw output matching - + # Check shapes assert outputs1.shape == outputs2.shape, f"Output shape mismatch: {outputs1.shape} vs {outputs2.shape}" - + diff = torch.abs(outputs1 - outputs2) max_diff = diff.max().item() mean_diff = diff.mean().item() - + logger.info(f"Max diff: {max_diff}") logger.info(f"Mean diff: {mean_diff}") - + # Thresholds # Qwen2.5-VL RoPE is complex, if our implementation is slightly off (e.g. float32 conversion logic in RoPE), # differences might appear. But should be small. atol = 5e-2 if precision_str == "bf16" else 1e-3 # relaxed for bf16 - + if max_diff > atol: logger.warning(f"Max diff {max_diff} > {atol}. Checking if it's acceptable...") # If mean diff is small, maybe just outliers assert mean_diff < atol, f"Mean diff {mean_diff} too high" else: logger.info("Outputs match within tolerance.") - diff --git a/fastvideo/tests/encoders/test_siglip_encoder.py b/fastvideo/tests/encoders/test_siglip_encoder.py index 60c1adb5a..b308d5ded 100644 --- a/fastvideo/tests/encoders/test_siglip_encoder.py +++ b/fastvideo/tests/encoders/test_siglip_encoder.py @@ -38,25 +38,21 @@ def test_siglip_encoder(): - Produce nearly identical outputs for the same input """ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - + logger.info("Loading SigLIP models from %s/%s", MODEL_ID, IMAGE_ENCODER_SUBFOLDER) - + # Load HuggingFace implementation - hf_model = HFSiglipVisionModel.from_pretrained( - MODEL_ID, subfolder=IMAGE_ENCODER_SUBFOLDER - ).to(torch.float16).to(device).eval() - + hf_model = HFSiglipVisionModel.from_pretrained(MODEL_ID, subfolder=IMAGE_ENCODER_SUBFOLDER).to( + torch.float16).to(device).eval() + # Load the config from Hugging Face and extract vision_config - config_path = hf_hub_download( - repo_id=MODEL_ID, - filename=f"{IMAGE_ENCODER_SUBFOLDER}/config.json" - ) + config_path = hf_hub_download(repo_id=MODEL_ID, filename=f"{IMAGE_ENCODER_SUBFOLDER}/config.json") with open(config_path) as f: full_config = json.load(f) - + # Get vision config from the full config vision_config_dict = full_config.get("vision_config", full_config) - + # Create FastVideo config with vision-specific settings arch_config = SiglipVisionArchConfig( hidden_size=vision_config_dict.get("hidden_size", 1152), @@ -66,36 +62,29 @@ def test_siglip_encoder(): num_hidden_layers=vision_config_dict.get("num_hidden_layers", 27), patch_size=vision_config_dict.get("patch_size", 14), ) - + config = SiglipVisionConfig(arch_config=arch_config) - + # Create FastVideo model from fastvideo.models.encoders.siglip import SiglipVisionModel fv_model = SiglipVisionModel(config).to(torch.float16).to(device) - + # Load weights from safetensors via Hugging Face - weights_path = hf_hub_download( - repo_id=MODEL_ID, - filename=f"{IMAGE_ENCODER_SUBFOLDER}/model.safetensors" - ) + weights_path = hf_hub_download(repo_id=MODEL_ID, filename=f"{IMAGE_ENCODER_SUBFOLDER}/model.safetensors") state_dict = load_file(weights_path) # Filter to only vision_model weights (keep the vision_model. prefix) - vision_weights = [ - (name, weight) - for name, weight in state_dict.items() - if name.startswith("vision_model.") - ] + vision_weights = [(name, weight) for name, weight in state_dict.items() if name.startswith("vision_model.")] fv_model.load_weights(vision_weights) fv_model.eval() - + # Sanity check weights between the two models logger.info("Comparing model weights for sanity check...") params1 = dict(hf_model.named_parameters()) params2 = dict(fv_model.named_parameters()) - + logger.info("HuggingFace model has %d parameters", len(params1)) logger.info("FastVideo model has %d parameters", len(params2)) - + # Compare non-stacked parameters # Note: HF uses param names directly, FV adds "vision_model." prefix for name1, param1 in sorted(params1.items()): @@ -113,41 +102,41 @@ def test_siglip_encoder(): param2 = params2[name2] param2 = param2.to_local().to(device) if isinstance(param2, DTensor) else param2.to(device) assert_close(param1, param2, atol=1e-4, rtol=1e-4) - + gc.collect() torch.cuda.empty_cache() - + # Test with sample images batch_size = 2 image_size = fv_model.config.arch_config.image_size - + # Create random pixel values pixel_values = torch.randn(batch_size, 3, image_size, image_size).to(device).to(torch.float16) - + logger.info("Testing SigLIP encoder with random pixel values of shape %s", pixel_values.shape) - + with torch.no_grad(): # Get embeddings from HuggingFace implementation hf_outputs = hf_model(pixel_values=pixel_values) - + # Get embeddings from FastVideo implementation with set_forward_context(current_timestep=0, attn_metadata=None): fv_outputs = fv_model(pixel_values=pixel_values) - + # Compare last hidden states hf_hidden_state = hf_outputs.last_hidden_state fv_hidden_state = fv_outputs.last_hidden_state - + logger.info("HF hidden state shape: %s", hf_hidden_state.shape) logger.info("FV hidden state shape: %s", fv_hidden_state.shape) - + assert hf_hidden_state.shape == fv_hidden_state.shape, \ f"Hidden state shapes don't match: {hf_hidden_state.shape} vs {fv_hidden_state.shape}" - + # Compare outputs with tolerance for numerical differences assert_close(hf_hidden_state, fv_hidden_state, atol=1e-2, rtol=1e-3) - + logger.info("SigLIP encoder test passed - outputs match within tolerance") - + gc.collect() torch.cuda.empty_cache() diff --git a/fastvideo/tests/encoders/test_t5_encoder.py b/fastvideo/tests/encoders/test_t5_encoder.py index bdaa683d4..8d69d8cec 100644 --- a/fastvideo/tests/encoders/test_t5_encoder.py +++ b/fastvideo/tests/encoders/test_t5_encoder.py @@ -26,9 +26,7 @@ os.environ["MASTER_PORT"] = "29503" @pytest.fixture def t5_model_paths_and_config(): base_model_path = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" - model_path = maybe_download_model(base_model_path, - local_dir=os.path.join( - 'data', base_model_path)) + model_path = maybe_download_model(base_model_path, local_dir=os.path.join('data', base_model_path)) text_encoder_path = os.path.join(model_path, "text_encoder") tokenizer_path = os.path.join(model_path, "tokenizer") return text_encoder_path, tokenizer_path, WanT2V480PConfig() @@ -38,9 +36,7 @@ def t5_model_paths_and_config(): def t5_large_model_paths_and_config(): base_model_path = "nvidia/Cosmos-Predict2-2B-Video2World" local_dir = os.path.join('data', base_model_path) - skip_if_gated_repo_inaccessible(base_model_path, - local_path=local_dir, - test_name="Cosmos T5-large encoder test") + skip_if_gated_repo_inaccessible(base_model_path, local_path=local_dir, test_name="Cosmos T5-large encoder test") model_path = maybe_download_model(base_model_path, local_dir=local_dir) text_encoder_path = os.path.join(model_path, "text_encoder") tokenizer_path = os.path.join(model_path, "tokenizer") @@ -57,14 +53,10 @@ def test_t5_encoder(t5_model_paths_and_config): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") precision_str = "fp32" precision = PRECISION_TO_TYPE[precision_str] - model1 = UMT5EncoderModel.from_pretrained(text_encoder_path).to( - precision).to(device).eval() + model1 = UMT5EncoderModel.from_pretrained(text_encoder_path).to(precision).to(device).eval() tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - - args = FastVideoArgs(model_path=text_encoder_path, - pipeline_config=pipeline_config, - pin_cpu_memory=False) + args = FastVideoArgs(model_path=text_encoder_path, pipeline_config=pipeline_config, pin_cpu_memory=False) loader = TextEncoderLoader() model2 = loader.load(text_encoder_path, args) model2 = model2.to(precision) @@ -84,7 +76,7 @@ def test_t5_encoder(t5_model_paths_and_config): "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_0.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_1.weight",\ "encoder.block.{}.layer.1.DenseReluDense.wo.weight", \ "encoder.block.{}.layer.1.layer_norm.weight", "encoder.final_layer_norm.weight"] - + for idx in range(hf_config.num_hidden_layers): for w in weights: name1 = w.format(idx) @@ -93,13 +85,9 @@ def test_t5_encoder(t5_model_paths_and_config): p2 = params2[name2] p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1) assert_close(p1, p2, atol=1e-4, rtol=1e-4) - # Test with some sample prompts - prompts = [ - "Once upon a time", "The quick brown fox jumps over", - "In a galaxy far, far away" - ] + prompts = ["Once upon a time", "The quick brown fox jumps over", "In a galaxy far, far away"] logger.info("Testing T5 encoder with sample prompts") @@ -108,10 +96,7 @@ def test_t5_encoder(t5_model_paths_and_config): logger.info("Testing prompt: %s", prompt) # Tokenize the prompt - tokens = tokenizer(prompt, - padding="max_length", - max_length=512, - truncation=True, + tokens = tokenizer(prompt, padding="max_length", max_length=512, truncation=True, return_tensors="pt").to(device) # Get outputs from HuggingFace implementation @@ -148,13 +133,10 @@ def test_t5_large_encoder(t5_large_model_paths_and_config): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") precision_str = "fp32" precision = PRECISION_TO_TYPE[precision_str] - model1 = T5EncoderModel.from_pretrained(text_encoder_path).to( - precision).to(device).eval() + model1 = T5EncoderModel.from_pretrained(text_encoder_path).to(precision).to(device).eval() tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - args = FastVideoArgs(model_path=text_encoder_path, - pipeline_config=pipeline_config, - pin_cpu_memory=False) + args = FastVideoArgs(model_path=text_encoder_path, pipeline_config=pipeline_config, pin_cpu_memory=False) loader = TextEncoderLoader() model2 = loader.load(text_encoder_path, args) model2 = model2.to(precision) @@ -168,12 +150,12 @@ def test_t5_large_encoder(t5_large_model_paths_and_config): # Check number of parameters logger.info("Model1 has %s parameters", len(params1)) logger.info("Model2 has %s parameters", len(params2)) - + # Print parameter names for comparison logger.info("Model1 parameters:") for name in sorted(params1.keys()): logger.info(" %s: %s", name, params1[name].shape) - + logger.info("Model2 parameters:") for name in sorted(params2.keys()): logger.info(" %s: %s", name, params2[name].shape) @@ -182,7 +164,7 @@ def test_t5_large_encoder(t5_large_model_paths_and_config): weights = ["encoder.block.{}.layer.0.layer_norm.weight", "encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight", \ "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi.weight", \ "encoder.block.{}.layer.1.DenseReluDense.wo.weight", "encoder.final_layer_norm.weight"] - + for idx in range(hf_config.num_hidden_layers): for w in weights: name1 = w.format(idx) @@ -191,13 +173,9 @@ def test_t5_large_encoder(t5_large_model_paths_and_config): p2 = params2[name2] p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1) assert_close(p1, p2, atol=1e-4, rtol=1e-4) - # Test with some sample prompts - prompts = [ - "Once upon a time", "The quick brown fox jumps over", - "In a galaxy far, far away" - ] + prompts = ["Once upon a time", "The quick brown fox jumps over", "In a galaxy far, far away"] logger.info("Testing T5 Large encoder with sample prompts") @@ -206,10 +184,7 @@ def test_t5_large_encoder(t5_large_model_paths_and_config): logger.info("Testing prompt: %s", prompt) # Tokenize the prompt - tokens = tokenizer(prompt, - padding="max_length", - max_length=512, - truncation=True, + tokens = tokenizer(prompt, padding="max_length", max_length=512, truncation=True, return_tensors="pt").to(device) # Get outputs from HuggingFace implementation diff --git a/fastvideo/tests/entrypoints/streaming/test_auxiliaries.py b/fastvideo/tests/entrypoints/streaming/test_auxiliaries.py index 3be3ef47c..758c4d0bf 100644 --- a/fastvideo/tests/entrypoints/streaming/test_auxiliaries.py +++ b/fastvideo/tests/entrypoints/streaming/test_auxiliaries.py @@ -34,7 +34,6 @@ from fastvideo.entrypoints.streaming.session_logger import ( SessionLogger, ) - # ---------------------------------------------------------------------- # Safety # ---------------------------------------------------------------------- @@ -60,8 +59,7 @@ class TestPromptSafetyFilter: assert result.decision is SafetyDecision.UNAVAILABLE def test_block_when_classifier_flags_unsafe(self, monkeypatch, tmp_path): - fake_model = types.SimpleNamespace( - predict=lambda text, k=1: (["__label__unsafe"], [0.95])) + fake_model = types.SimpleNamespace(predict=lambda text, k=1: (["__label__unsafe"], [0.95])) stub = types.SimpleNamespace(load_model=lambda _p: fake_model) monkeypatch.setitem(sys.modules, "fasttext", stub) model_path = str(tmp_path / "m.bin") @@ -73,8 +71,7 @@ class TestPromptSafetyFilter: assert result.score == pytest.approx(0.95) def test_allow_when_classifier_flags_safe(self, monkeypatch, tmp_path): - fake_model = types.SimpleNamespace( - predict=lambda text, k=1: (["__label__safe"], [0.99])) + fake_model = types.SimpleNamespace(predict=lambda text, k=1: (["__label__safe"], [0.99])) stub = types.SimpleNamespace(load_model=lambda _p: fake_model) monkeypatch.setitem(sys.modules, "fasttext", stub) f = PromptSafetyFilter(classifier_path="ignored", enabled=True) @@ -82,12 +79,10 @@ class TestPromptSafetyFilter: assert result.decision is SafetyDecision.ALLOW def test_below_threshold_allows_even_if_unsafe_label(self, monkeypatch): - fake_model = types.SimpleNamespace( - predict=lambda text, k=1: (["__label__unsafe"], [0.3])) + fake_model = types.SimpleNamespace(predict=lambda text, k=1: (["__label__unsafe"], [0.3])) stub = types.SimpleNamespace(load_model=lambda _p: fake_model) monkeypatch.setitem(sys.modules, "fasttext", stub) - f = PromptSafetyFilter( - classifier_path="m", enabled=True, block_threshold=0.5) + f = PromptSafetyFilter(classifier_path="m", enabled=True, block_threshold=0.5) assert f.classify("x").decision is SafetyDecision.ALLOW def test_first_blocked_returns_first_hit(self, monkeypatch): @@ -96,8 +91,7 @@ class TestPromptSafetyFilter: (["__label__unsafe"], [0.9]), (["__label__safe"], [0.9]), ]) - fake_model = types.SimpleNamespace( - predict=lambda text, k=1: next(responses)) + fake_model = types.SimpleNamespace(predict=lambda text, k=1: next(responses)) stub = types.SimpleNamespace(load_model=lambda _p: fake_model) monkeypatch.setitem(sys.modules, "fasttext", stub) f = PromptSafetyFilter(classifier_path="m", enabled=True) @@ -115,17 +109,22 @@ class TestRewriteSplit: def test_plain_lines(self): assert _split_response("one\ntwo\nthree", limit=3) == [ - "one", "two", "three", + "one", + "two", + "three", ] def test_numbered_list(self): assert _split_response("1. first\n2. second", limit=3) == [ - "first", "second", + "first", + "second", ] def test_bulleted_list(self): assert _split_response("- one\n* two\n• three", limit=3) == [ - "one", "two", "three", + "one", + "two", + "three", ] def test_respects_limit(self): @@ -159,9 +158,7 @@ class TestBuildRewrite: def test_returns_limited_alternatives(self): import asyncio - result = asyncio.run(build_rewrite( - _StubEnhancer(), "seed", - options=RewriteOptions(count=2))) + result = asyncio.run(build_rewrite(_StubEnhancer(), "seed", options=RewriteOptions(count=2))) assert result.seed_prompt == "seed" assert result.alternatives == ["alpha", "beta"] assert result.provider == "stub" @@ -229,11 +226,11 @@ class TestMockServer: from fastvideo.entrypoints.streaming.mock_server import MockGenerator gen = MockGenerator() - result = gen.generate(GenerationRequest( - prompt="x", - sampling=SamplingConfig( - num_frames=3, height=32, width=32, num_inference_steps=1), - )) + result = gen.generate( + GenerationRequest( + prompt="x", + sampling=SamplingConfig(num_frames=3, height=32, width=32, num_inference_steps=1), + )) assert len(result["frames"]) == 3 assert result["frames"][0].shape == (32, 32, 3) assert result["state"].kind == "ltx2.v1" diff --git a/fastvideo/tests/entrypoints/streaming/test_gpu_pool.py b/fastvideo/tests/entrypoints/streaming/test_gpu_pool.py index 10ce06af6..2f02d6843 100644 --- a/fastvideo/tests/entrypoints/streaming/test_gpu_pool.py +++ b/fastvideo/tests/entrypoints/streaming/test_gpu_pool.py @@ -31,7 +31,6 @@ from fastvideo.entrypoints.streaming.gpu_pool import ( _WorkerHandle, ) - # ---------------------------------------------------------------------- # In-process pool # ---------------------------------------------------------------------- @@ -53,8 +52,7 @@ class _MockGenerator: class TestInProcessGpuPool: def test_is_gpu_pool(self): - assert isinstance( - InProcessGpuPool(_MockGenerator()), GpuPool) + assert isinstance(InProcessGpuPool(_MockGenerator()), GpuPool) def test_acquire_returns_deterministic_assignment(self): pool = InProcessGpuPool(_MockGenerator(), gpu_id=7) @@ -83,8 +81,7 @@ class TestInProcessGpuPool: async def run(): with pytest.raises(RuntimeError): - await pool.run( - "sess-a", GenerationRequest(prompt="hi")) + await pool.run("sess-a", GenerationRequest(prompt="hi")) asyncio.run(run()) @@ -93,8 +90,7 @@ class TestInProcessGpuPool: async def run(): await pool.acquire("sess-a") - return await pool.run( - "sess-a", GenerationRequest(prompt="hi")) + return await pool.run("sess-a", GenerationRequest(prompt="hi")) result = asyncio.run(run()) assert result["prompt_echo"] == "hi" @@ -106,8 +102,7 @@ class TestInProcessGpuPool: await pool.acquire("sess-a") await pool.release("sess-a") with pytest.raises(RuntimeError): - await pool.run( - "sess-a", GenerationRequest(prompt="hi")) + await pool.run("sess-a", GenerationRequest(prompt="hi")) asyncio.run(run()) @@ -234,6 +229,7 @@ def _thread_worker_factory(generator_builder): threading.Thread(target=_await_ready, daemon=True).start() class _FakeProcess: + def __init__(self, stop: threading.Event, worker: _ThreadWorker): self._stop = stop self._worker = worker @@ -273,8 +269,7 @@ def pool_factory(): generator_config=GeneratorConfig(model_path="/models/fake"), pool_config=GpuPoolConfig(num_workers=num_workers), warmup_config=WarmupConfig(enabled=False), - worker_factory=_thread_worker_factory( - lambda gpu_id: _MockGenerator()), + worker_factory=_thread_worker_factory(lambda gpu_id: _MockGenerator()), ) await pool.start() return pool @@ -285,6 +280,7 @@ def pool_factory(): class TestSubprocessGpuPool: def test_start_spawns_requested_workers(self, pool_factory): + async def run(): pool = await pool_factory(num_workers=3) try: @@ -297,6 +293,7 @@ class TestSubprocessGpuPool: asyncio.run(run()) def test_acquire_decrements_available(self, pool_factory): + async def run(): pool = await pool_factory(num_workers=2) try: @@ -309,6 +306,7 @@ class TestSubprocessGpuPool: asyncio.run(run()) def test_acquire_timeout_when_all_busy(self, pool_factory): + async def run(): pool = await pool_factory(num_workers=1) try: @@ -321,12 +319,12 @@ class TestSubprocessGpuPool: asyncio.run(run()) def test_run_returns_worker_result(self, pool_factory): + async def run(): pool = await pool_factory(num_workers=1) try: await pool.acquire("sess-a") - result = await pool.run( - "sess-a", GenerationRequest(prompt="hello")) + result = await pool.run("sess-a", GenerationRequest(prompt="hello")) assert result["prompt_echo"] == "hello" finally: await pool.shutdown() @@ -334,6 +332,7 @@ class TestSubprocessGpuPool: asyncio.run(run()) def test_release_returns_worker_to_pool(self, pool_factory): + async def run(): pool = await pool_factory(num_workers=1) try: @@ -348,6 +347,7 @@ class TestSubprocessGpuPool: asyncio.run(run()) def test_sticky_binding_across_multiple_runs(self, pool_factory): + async def run(): pool = await pool_factory(num_workers=2) try: @@ -363,18 +363,19 @@ class TestSubprocessGpuPool: asyncio.run(run()) def test_run_without_acquire_raises(self, pool_factory): + async def run(): pool = await pool_factory(num_workers=1) try: with pytest.raises(RuntimeError): - await pool.run( - "sess-x", GenerationRequest(prompt="x")) + await pool.run("sess-x", GenerationRequest(prompt="x")) finally: await pool.shutdown() asyncio.run(run()) def test_shutdown_is_idempotent(self, pool_factory): + async def run(): pool = await pool_factory(num_workers=2) await pool.shutdown() @@ -404,10 +405,13 @@ class TestSubprocessGpuPoolFailureModes: boot_ok.set() class _AliveProcess: + def is_alive(self) -> bool: return True + def join(self, timeout: float | None = None) -> None: return + def kill(self) -> None: return diff --git a/fastvideo/tests/entrypoints/streaming/test_health_routes.py b/fastvideo/tests/entrypoints/streaming/test_health_routes.py index c385bd394..2516fde7f 100644 --- a/fastvideo/tests/entrypoints/streaming/test_health_routes.py +++ b/fastvideo/tests/entrypoints/streaming/test_health_routes.py @@ -12,24 +12,16 @@ from fastapi import FastAPI # noqa: E402 from starlette.testclient import TestClient # noqa: E402 from fastvideo.api.schema import ( # noqa: E402 - GeneratorConfig, - GenerationRequest, - SamplingConfig, - ServeConfig, - StreamingConfig, + GeneratorConfig, GenerationRequest, SamplingConfig, ServeConfig, StreamingConfig, ) from fastvideo.entrypoints.streaming.gpu_pool import ( # noqa: E402 - GpuPool, - PoolAssignment, - PoolHealth, + GpuPool, PoolAssignment, PoolHealth, ) from fastvideo.entrypoints.streaming.health import ( # noqa: E402 - build_health_router, - get_pool_status, + build_health_router, get_pool_status, ) from fastvideo.entrypoints.streaming.server import build_app # noqa: E402 - _STATUS_KEYS = { "total_gpus", "available_gpus", @@ -88,7 +80,9 @@ def _fake_status() -> dict[str, Any]: "warmup_enabled": True, "warmup_success": True, "warmup_error": None, - "warmup_timings": {"load": 1.2}, + "warmup_timings": { + "load": 1.2 + }, }, 1: { "ready": False, @@ -114,15 +108,13 @@ def _build_health_client(pool=None) -> TestClient: def _build_serve_config() -> ServeConfig: return ServeConfig( generator=GeneratorConfig(model_path="/models/fake"), - default_request=GenerationRequest( - sampling=SamplingConfig( - num_frames=12, - height=64, - width=64, - fps=12, - num_inference_steps=1, - ), - ), + default_request=GenerationRequest(sampling=SamplingConfig( + num_frames=12, + height=64, + width=64, + fps=12, + num_inference_steps=1, + ), ), streaming=StreamingConfig( session_timeout_seconds=60, generation_segment_cap=2, diff --git a/fastvideo/tests/entrypoints/streaming/test_prompt_enhancer.py b/fastvideo/tests/entrypoints/streaming/test_prompt_enhancer.py index e9bb5bc7c..f599e92ce 100644 --- a/fastvideo/tests/entrypoints/streaming/test_prompt_enhancer.py +++ b/fastvideo/tests/entrypoints/streaming/test_prompt_enhancer.py @@ -125,8 +125,7 @@ class TestAutoExtendAndRewrite: name = "cap" async def complete(self, request: LLMRequest) -> LLMResponse: - user = next(m.content for m in request.messages - if m.role == "user") + user = next(m.content for m in request.messages if m.role == "user") captured.append(user) return LLMResponse( content="next", @@ -146,8 +145,7 @@ class TestAutoExtendAndRewrite: name = "cap" async def complete(self, request: LLMRequest) -> LLMResponse: - user = next(m.content for m in request.messages - if m.role == "user") + user = next(m.content for m in request.messages if m.role == "user") captured.append(user) return LLMResponse( content="one\ntwo\nthree", @@ -177,8 +175,7 @@ class TestRegisterProvider: providers=[_StaticProvider("old", "a")], model="m", ) - enh.register_provider( - _StaticProvider("new-primary", "b"), priority=0) + enh.register_provider(_StaticProvider("new-primary", "b"), priority=0) assert enh.providers[0].name == "new-primary" def test_registered_provider_is_used_in_enhance(self): diff --git a/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py b/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py index a490e231e..66288ebef 100644 --- a/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py +++ b/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py @@ -102,8 +102,7 @@ class TestCerebrasProvider: monkeypatch.delenv("CEREBRAS_API_KEY", raising=False) provider = CerebrasProvider() with pytest.raises(LLMProviderError, match="CEREBRAS_API_KEY"): - asyncio.run(provider.complete( - LLMRequest(messages=[], model="m"))) + asyncio.run(provider.complete(LLMRequest(messages=[], model="m"))) def test_api_key_from_env(self, monkeypatch): monkeypatch.setenv("CEREBRAS_API_KEY", "from-env") @@ -121,18 +120,23 @@ class TestCerebrasProvider: response=_FakeResponse( status_code=200, payload={ - "choices": [{"message": {"content": "enhanced"}}], + "choices": [{ + "message": { + "content": "enhanced" + } + }], }, ), ) provider = CerebrasProvider(api_key="secret") - result = asyncio.run(provider.complete( - LLMRequest( - messages=[LLMMessage(role="user", content="a fox")], - model="gpt-oss-120b", - max_tokens=64, - temperature=0.7, - ))) + result = asyncio.run( + provider.complete( + LLMRequest( + messages=[LLMMessage(role="user", content="a fox")], + model="gpt-oss-120b", + max_tokens=64, + temperature=0.7, + ))) assert result.content == "enhanced" assert result.provider == "cerebras" assert result.model == "gpt-oss-120b" @@ -155,14 +159,12 @@ class TestCerebrasProvider: stub.HTTPError = _HTTPError stub.TimeoutException = _TimeoutException - stub.AsyncClient = lambda *a, **k: _FakeAsyncClient( - _HTTPError("boom"), captured=[]) + stub.AsyncClient = lambda *a, **k: _FakeAsyncClient(_HTTPError("boom"), captured=[]) monkeypatch.setitem(sys.modules, "httpx", stub) provider = CerebrasProvider(api_key="k") with pytest.raises(LLMProviderError, match="HTTP"): - asyncio.run(provider.complete( - LLMRequest(messages=[], model="m"))) + asyncio.run(provider.complete(LLMRequest(messages=[], model="m"))) def test_timeout_wrapped(self, monkeypatch): _install_fake_httpx(monkeypatch) @@ -175,8 +177,7 @@ class TestCerebrasProvider: stub.AsyncClient = raising_factory # type: ignore[attr-defined] provider = CerebrasProvider(api_key="k") with pytest.raises(LLMTimeoutError): - asyncio.run(provider.complete( - LLMRequest(messages=[], model="m"))) + asyncio.run(provider.complete(LLMRequest(messages=[], model="m"))) def test_4xx_raises_non_retryable_provider_error(self, monkeypatch): _install_fake_httpx( @@ -189,8 +190,7 @@ class TestCerebrasProvider: ) provider = CerebrasProvider(api_key="k") with pytest.raises(LLMProviderError, match="401") as excinfo: - asyncio.run(provider.complete( - LLMRequest(messages=[], model="m"))) + asyncio.run(provider.complete(LLMRequest(messages=[], model="m"))) assert excinfo.value.retryable is False def test_429_raises_retryable_provider_error(self, monkeypatch): @@ -204,8 +204,7 @@ class TestCerebrasProvider: ) provider = CerebrasProvider(api_key="k") with pytest.raises(LLMProviderError, match="429") as excinfo: - asyncio.run(provider.complete( - LLMRequest(messages=[], model="m"))) + asyncio.run(provider.complete(LLMRequest(messages=[], model="m"))) assert excinfo.value.retryable is True def test_5xx_raises_retryable_provider_error(self, monkeypatch): @@ -219,8 +218,7 @@ class TestCerebrasProvider: ) provider = CerebrasProvider(api_key="k") with pytest.raises(LLMProviderError, match="503") as excinfo: - asyncio.run(provider.complete( - LLMRequest(messages=[], model="m"))) + asyncio.run(provider.complete(LLMRequest(messages=[], model="m"))) assert excinfo.value.retryable is True def test_non_json_body_raises_provider_error(self, monkeypatch): @@ -240,19 +238,16 @@ class TestCerebrasProvider: ) provider = CerebrasProvider(api_key="k") with pytest.raises(LLMProviderError, match="non-JSON"): - asyncio.run(provider.complete( - LLMRequest(messages=[], model="m"))) + asyncio.run(provider.complete(LLMRequest(messages=[], model="m"))) def test_empty_choices_raises(self, monkeypatch): _install_fake_httpx( monkeypatch, - response=_FakeResponse( - status_code=200, payload={"choices": []}), + response=_FakeResponse(status_code=200, payload={"choices": []}), ) provider = CerebrasProvider(api_key="k") with pytest.raises(LLMProviderError, match="no choices"): - asyncio.run(provider.complete( - LLMRequest(messages=[], model="m"))) + asyncio.run(provider.complete(LLMRequest(messages=[], model="m"))) # ---------------------------------------------------------------------- @@ -269,8 +264,7 @@ class TestGroqProvider: monkeypatch.delenv("GROQ_API_KEY", raising=False) provider = GroqProvider() with pytest.raises(LLMProviderError, match="GROQ_API_KEY"): - asyncio.run(provider.complete( - LLMRequest(messages=[], model="m"))) + asyncio.run(provider.complete(LLMRequest(messages=[], model="m"))) def test_api_key_from_env(self, monkeypatch): monkeypatch.setenv("GROQ_API_KEY", "from-env") @@ -283,13 +277,17 @@ class TestGroqProvider: response=_FakeResponse( status_code=200, payload={ - "choices": [{"message": {"content": "groq out"}}], + "choices": [{ + "message": { + "content": "groq out" + } + }], }, ), ) provider = GroqProvider(api_key="secret") - result = asyncio.run(provider.complete( - LLMRequest( + result = asyncio.run( + provider.complete(LLMRequest( messages=[LLMMessage(role="user", content="a deer")], model="llama-3.1-70b", ))) diff --git a/fastvideo/tests/entrypoints/streaming/test_protocol.py b/fastvideo/tests/entrypoints/streaming/test_protocol.py index 5b55d517f..26ea5c5a5 100644 --- a/fastvideo/tests/entrypoints/streaming/test_protocol.py +++ b/fastvideo/tests/entrypoints/streaming/test_protocol.py @@ -53,7 +53,10 @@ class TestClientMessageParsing: "stream_mode": "av_fmp4", "continuation_state": { "kind": "ltx2.v1", - "payload": {"schema_version": 1, "segment_index": 2}, + "payload": { + "schema_version": 1, + "segment_index": 2 + }, }, } parsed = parse_client_message(raw) @@ -107,7 +110,10 @@ class TestServerMessageSerialization: def test_ltx2_stream_start(self): msg = Ltx2StreamStart( preset="ltx2_two_stage", - width=1024, height=1536, fps=24, num_frames=121, + width=1024, + height=1536, + fps=24, + num_frames=121, ) dumped = msg.model_dump() assert dumped["type"] == "ltx2_stream_start" @@ -133,7 +139,9 @@ class TestServerMessageSerialization: def test_media_segment_complete(self): msg = MediaSegmentComplete( - segment_idx=0, stream_id="abc", chunks=4, + segment_idx=0, + stream_id="abc", + chunks=4, ) dumped = msg.model_dump() assert dumped["chunks"] == 4 @@ -149,6 +157,8 @@ class TestServerMessageSerialization: def test_continuation_state_snapshot(self): msg = ContinuationStateSnapshot(state={ "kind": "ltx2.v1", - "payload": {"schema_version": 1}, + "payload": { + "schema_version": 1 + }, }) assert msg.model_dump()["state"]["kind"] == "ltx2.v1" diff --git a/fastvideo/tests/entrypoints/streaming/test_router.py b/fastvideo/tests/entrypoints/streaming/test_router.py index a1e92789d..1de5e5339 100644 --- a/fastvideo/tests/entrypoints/streaming/test_router.py +++ b/fastvideo/tests/entrypoints/streaming/test_router.py @@ -27,19 +27,14 @@ def _registry( num_primary: int = 1, num_secondary: int = 1, ) -> ReplicaRegistry: - replicas = [ - ReplicaEndpoint( - url=f"http://primary-{i}:8000", - primary=True, - ) - for i in range(num_primary) - ] + [ - ReplicaEndpoint( - url=f"http://secondary-{i}:8000", - primary=False, - ) - for i in range(num_secondary) - ] + replicas = [ReplicaEndpoint( + url=f"http://primary-{i}:8000", + primary=True, + ) for i in range(num_primary) + ] + [ReplicaEndpoint( + url=f"http://secondary-{i}:8000", + primary=False, + ) for i in range(num_secondary)] return ReplicaRegistry(replicas) @@ -57,12 +52,10 @@ class TestReplicaRegistry: async def promote(): primary = reg.primaries()[0] - await reg.record_success( - primary, recovery_threshold=1, latency_ms=1.0) + await reg.record_success(primary, recovery_threshold=1, latency_ms=1.0) # Also mark secondary healthy — primary should still win. sec = next(r for r in reg.all() if not r.primary) - await reg.record_success( - sec, recovery_threshold=1, latency_ms=1.0) + await reg.record_success(sec, recovery_threshold=1, latency_ms=1.0) return reg.select() pick = asyncio.run(promote()) @@ -77,10 +70,8 @@ class TestReplicaRegistry: sec = next(r for r in reg.all() if not r.primary) # Fail primary past threshold, succeed secondary. for _ in range(3): - await reg.record_failure( - primary, failure_threshold=3, reason="mock") - await reg.record_success( - sec, recovery_threshold=1, latency_ms=1.0) + await reg.record_failure(primary, failure_threshold=3, reason="mock") + await reg.record_success(sec, recovery_threshold=1, latency_ms=1.0) return reg.select() pick = asyncio.run(run()) @@ -93,11 +84,9 @@ class TestReplicaRegistry: async def run(): for _ in range(2): - await reg.record_failure( - primary, failure_threshold=3, reason="x") + await reg.record_failure(primary, failure_threshold=3, reason="x") assert primary.health.status is not ReplicaStatus.UNHEALTHY - await reg.record_failure( - primary, failure_threshold=3, reason="x") + await reg.record_failure(primary, failure_threshold=3, reason="x") return primary result = asyncio.run(run()) @@ -110,12 +99,10 @@ class TestReplicaRegistry: async def run(): for _ in range(3): - await reg.record_failure( - primary, failure_threshold=3, reason="x") + await reg.record_failure(primary, failure_threshold=3, reason="x") assert primary.health.status is ReplicaStatus.UNHEALTHY for _ in range(2): - await reg.record_success( - primary, recovery_threshold=2, latency_ms=5.0) + await reg.record_success(primary, recovery_threshold=2, latency_ms=5.0) return primary result = asyncio.run(run()) @@ -127,10 +114,8 @@ class TestReplicaRegistry: primary = reg.primaries()[0] async def run(): - await reg.record_failure( - primary, failure_threshold=10, reason="x") - await reg.record_success( - primary, recovery_threshold=1, latency_ms=1.0) + await reg.record_failure(primary, failure_threshold=10, reason="x") + await reg.record_success(primary, recovery_threshold=1, latency_ms=1.0) return primary result = asyncio.run(run()) @@ -161,10 +146,13 @@ class TestHealthCheckLoop: return 0.0, "mock failure" async def run() -> None: - task = asyncio.create_task(run_health_check_loop( - registry=reg, config=config, stop_event=stop_event, - http_get=probe, - )) + task = asyncio.create_task( + run_health_check_loop( + registry=reg, + config=config, + stop_event=stop_event, + http_get=probe, + )) await asyncio.sleep(0.05) stop_event.set() await task @@ -246,7 +234,7 @@ class TestUnknownToHealthyImmediate: await reg.record_success(primary, recovery_threshold=2, latency_ms=1.0) assert primary.health.status is ReplicaStatus.UNHEALTHY # 1/2 successes await reg.record_success(primary, recovery_threshold=2, latency_ms=1.0) - assert primary.health.status is ReplicaStatus.HEALTHY # 2/2 successes + assert primary.health.status is ReplicaStatus.HEALTHY # 2/2 successes asyncio.run(run()) diff --git a/fastvideo/tests/entrypoints/streaming/test_server.py b/fastvideo/tests/entrypoints/streaming/test_server.py index d878ea4f9..15c38b186 100644 --- a/fastvideo/tests/entrypoints/streaming/test_server.py +++ b/fastvideo/tests/entrypoints/streaming/test_server.py @@ -17,16 +17,10 @@ pytest.importorskip("starlette") from starlette.testclient import TestClient # noqa: E402 from fastvideo.api.schema import ( # noqa: E402 - ContinuationState, - GeneratorConfig, - SamplingConfig, - ServeConfig, - StreamingConfig, - GenerationRequest, + ContinuationState, GeneratorConfig, SamplingConfig, ServeConfig, StreamingConfig, GenerationRequest, ) from fastvideo.entrypoints.streaming.server import build_app # noqa: E402 - _FFMPEG_AVAILABLE = shutil.which("ffmpeg") is not None @@ -39,10 +33,7 @@ class _MockGenerator: return_state: bool = True def generate(self, request: GenerationRequest) -> dict[str, Any]: - frames = [ - np.full((self.height, self.width, 3), i * 5, dtype=np.uint8) - for i in range(self.num_frames) - ] + frames = [np.full((self.height, self.width, 3), i * 5, dtype=np.uint8) for i in range(self.num_frames)] state = (ContinuationState( kind="ltx2.v1", payload={ @@ -61,15 +52,13 @@ class _MockGenerator: def _build_serve_config() -> ServeConfig: return ServeConfig( generator=GeneratorConfig(model_path="/models/fake"), - default_request=GenerationRequest( - sampling=SamplingConfig( - num_frames=12, - height=64, - width=64, - fps=12, - num_inference_steps=1, - ), - ), + default_request=GenerationRequest(sampling=SamplingConfig( + num_frames=12, + height=64, + width=64, + fps=12, + num_inference_steps=1, + ), ), streaming=StreamingConfig( session_timeout_seconds=60, generation_segment_cap=2, @@ -132,7 +121,10 @@ class TestSessionHandshake: "preset": "ltx2_two_stage", "continuation_state": { "kind": "ltx2.v1", - "payload": {"schema_version": 1, "segment_index": 3}, + "payload": { + "schema_version": 1, + "segment_index": 3 + }, }, }) # Drain handshake frames @@ -153,8 +145,7 @@ class TestSegmentFlow: def test_segment_generates_media_init_plus_complete(self): client, generator = _build_client() with client.websocket_connect("/v1/stream") as ws: - ws.send_json({"type": "session_init_v2", - "preset": "ltx2_two_stage"}) + ws.send_json({"type": "session_init_v2", "preset": "ltx2_two_stage"}) for _ in range(3): ws.receive_json() # queue_status + gpu_assigned + stream_start @@ -195,8 +186,7 @@ class TestContinuationStatePersistence: pytest.skip("ffmpeg not installed") client, generator = _build_client() with client.websocket_connect("/v1/stream") as ws: - ws.send_json({"type": "session_init_v2", - "preset": "ltx2_two_stage"}) + ws.send_json({"type": "session_init_v2", "preset": "ltx2_two_stage"}) for _ in range(3): ws.receive_json() ws.send_json({ diff --git a/fastvideo/tests/entrypoints/streaming/test_session.py b/fastvideo/tests/entrypoints/streaming/test_session.py index b9be7eae8..dd0748df1 100644 --- a/fastvideo/tests/entrypoints/streaming/test_session.py +++ b/fastvideo/tests/entrypoints/streaming/test_session.py @@ -78,23 +78,20 @@ class TestSessionStateMachine: class TestSessionManager: def test_create_assigns_unique_ids(self): - mgr = SessionManager( - segment_cap=4, session_timeout_seconds=60, max_sessions=2) + mgr = SessionManager(segment_cap=4, session_timeout_seconds=60, max_sessions=2) a = mgr.create() b = mgr.create() assert a.id != b.id assert len(mgr) == 2 def test_max_sessions_enforced(self): - mgr = SessionManager( - segment_cap=4, session_timeout_seconds=60, max_sessions=1) + mgr = SessionManager(segment_cap=4, session_timeout_seconds=60, max_sessions=1) mgr.create() with pytest.raises(SessionRejected): mgr.create() def test_close_releases_slot(self): - mgr = SessionManager( - segment_cap=4, session_timeout_seconds=60, max_sessions=1) + mgr = SessionManager(segment_cap=4, session_timeout_seconds=60, max_sessions=1) s = mgr.create() mgr.close(s.id) assert len(mgr) == 0 @@ -102,8 +99,7 @@ class TestSessionManager: mgr.create() def test_reap_timed_out_flags_idle_sessions(self): - mgr = SessionManager( - segment_cap=4, session_timeout_seconds=1, max_sessions=4) + mgr = SessionManager(segment_cap=4, session_timeout_seconds=1, max_sessions=4) s = mgr.create() s.transition(SessionState.QUEUED) s.transition(SessionState.GPU_BINDING) @@ -113,8 +109,7 @@ class TestSessionManager: assert s.id in dead def test_reap_skips_terminal_states(self): - mgr = SessionManager( - segment_cap=4, session_timeout_seconds=1, max_sessions=4) + mgr = SessionManager(segment_cap=4, session_timeout_seconds=1, max_sessions=4) s = mgr.create() s.transition(SessionState.QUEUED) s.transition(SessionState.ERROR) @@ -122,8 +117,7 @@ class TestSessionManager: assert s.id not in mgr.reap_timed_out() def test_active_sessions_filter(self): - mgr = SessionManager( - segment_cap=4, session_timeout_seconds=60, max_sessions=4) + mgr = SessionManager(segment_cap=4, session_timeout_seconds=60, max_sessions=4) a = mgr.create() a.transition(SessionState.QUEUED) a.transition(SessionState.GPU_BINDING) diff --git a/fastvideo/tests/entrypoints/streaming/test_session_init_image.py b/fastvideo/tests/entrypoints/streaming/test_session_init_image.py index 21faa67a9..b6d94dd14 100644 --- a/fastvideo/tests/entrypoints/streaming/test_session_init_image.py +++ b/fastvideo/tests/entrypoints/streaming/test_session_init_image.py @@ -10,8 +10,7 @@ import pytest from PIL import Image from fastvideo.entrypoints.streaming.session_init_image import ( - persist_session_init_image, -) + persist_session_init_image, ) def _png_bytes(size: tuple[int, int] = (64, 64)) -> bytes: @@ -32,11 +31,13 @@ class TestPersistSessionInitImage: def test_png_payload_persists(self, tmp_path): data = _png_bytes() - image = persist_session_init_image({ - "mime": "image/png", - "name": "ref.png", - "data": base64.b64encode(data).decode("ascii"), - }, output_dir=str(tmp_path)) + image = persist_session_init_image( + { + "mime": "image/png", + "name": "ref.png", + "data": base64.b64encode(data).decode("ascii"), + }, + output_dir=str(tmp_path)) assert image is not None assert os.path.exists(image.path) assert image.mime == "image/png" @@ -66,11 +67,13 @@ class TestPersistSessionInitImage: }, output_dir=str(tmp_path)) def test_display_name_sanitized(self, tmp_path): - image = persist_session_init_image({ - "mime": "image/png", - "name": "../evil/../name.png", - "data": base64.b64encode(_png_bytes()).decode("ascii"), - }, output_dir=str(tmp_path)) + image = persist_session_init_image( + { + "mime": "image/png", + "name": "../evil/../name.png", + "data": base64.b64encode(_png_bytes()).decode("ascii"), + }, + output_dir=str(tmp_path)) assert image is not None assert image.display_name == "name.png" @@ -81,10 +84,11 @@ class TestPersistSessionInitImage: mod._MAX_IMAGE_BYTES = 100 try: with pytest.raises(ValueError, match="limit"): - persist_session_init_image({ - "mime": "image/png", - "data": base64.b64encode(_png_bytes((512, 512))).decode( - "ascii"), - }, output_dir=str(tmp_path)) + persist_session_init_image( + { + "mime": "image/png", + "data": base64.b64encode(_png_bytes((512, 512))).decode("ascii"), + }, + output_dir=str(tmp_path)) finally: mod._MAX_IMAGE_BYTES = original diff --git a/fastvideo/tests/entrypoints/streaming/test_session_store.py b/fastvideo/tests/entrypoints/streaming/test_session_store.py index c5402fa91..55aeb9c0d 100644 --- a/fastvideo/tests/entrypoints/streaming/test_session_store.py +++ b/fastvideo/tests/entrypoints/streaming/test_session_store.py @@ -137,7 +137,7 @@ class TestInMemorySessionStore: ContinuationState(kind="ltx2.v1", payload={"i": i}), ) - threads = [threading.Thread(target=write, args=(i,)) for i in range(200)] + threads = [threading.Thread(target=write, args=(i, )) for i in range(200)] for t in threads: t.start() for t in threads: @@ -154,9 +154,7 @@ class TestSnapshotHydrateRoundTrip: typed = LTX2ContinuationState( segment_index=4, - video_frames=[ - np.full((32, 32, 3), i * 5, dtype=np.uint8) for i in range(3) - ], + video_frames=[np.full((32, 32, 3), i * 5, dtype=np.uint8) for i in range(3)], audio_latents=torch.randn(1, 4, 8, 32, dtype=torch.float32), audio_sample_rate=24000, audio_conditioning_num_frames=5, @@ -177,10 +175,8 @@ class TestSnapshotHydrateRoundTrip: rebuilt = sessions.snapshot(new_sid) assert rebuilt is snapshot - restored = LTX2ContinuationState.from_continuation_state( - rebuilt, blob_store=blob_store) + restored = LTX2ContinuationState.from_continuation_state(rebuilt, blob_store=blob_store) assert restored.segment_index == typed.segment_index assert restored.audio_sample_rate == typed.audio_sample_rate - torch.testing.assert_close( - restored.audio_latents, typed.audio_latents) + torch.testing.assert_close(restored.audio_latents, typed.audio_latents) assert len(restored.video_frames) == len(typed.video_frames) diff --git a/fastvideo/tests/entrypoints/streaming/test_stream_encoder.py b/fastvideo/tests/entrypoints/streaming/test_stream_encoder.py index 02b193cae..a48fce511 100644 --- a/fastvideo/tests/entrypoints/streaming/test_stream_encoder.py +++ b/fastvideo/tests/entrypoints/streaming/test_stream_encoder.py @@ -28,9 +28,9 @@ def _frame(width: int, height: int, value: int = 128) -> np.ndarray: def test_encoder_emits_init_then_media_chunks(): + async def run(): - enc = FragmentedMP4Encoder( - width=64, height=64, fps=24, segment_idx=0) + enc = FragmentedMP4Encoder(width=64, height=64, fps=24, segment_idx=0) chunks: list[FragmentedMP4Chunk] = [] async with enc: frames = [_frame(64, 64, v) for v in range(4, 28)] @@ -46,9 +46,9 @@ def test_encoder_emits_init_then_media_chunks(): def test_encoder_init_chunk_is_fmp4(): """The first chunk must contain the ``ftyp`` box (fMP4 init segment).""" + async def run(): - enc = FragmentedMP4Encoder( - width=64, height=64, fps=24, segment_idx=0) + enc = FragmentedMP4Encoder(width=64, height=64, fps=24, segment_idx=0) first_chunk = None async with enc: async for chunk in enc.encode([_frame(64, 64, 20)] * 24): @@ -64,9 +64,9 @@ def test_encoder_init_chunk_is_fmp4(): def test_encoder_rejects_non_ndarray_frames(): + async def run(): - enc = FragmentedMP4Encoder( - width=64, height=64, fps=24, segment_idx=0) + enc = FragmentedMP4Encoder(width=64, height=64, fps=24, segment_idx=0) async with enc: with pytest.raises(TypeError): async for _ in enc.encode(["not-a-frame"]): @@ -76,22 +76,21 @@ def test_encoder_rejects_non_ndarray_frames(): def test_encoder_rejects_wrong_shape(): + async def run(): - enc = FragmentedMP4Encoder( - width=64, height=64, fps=24, segment_idx=0) + enc = FragmentedMP4Encoder(width=64, height=64, fps=24, segment_idx=0) async with enc: with pytest.raises(ValueError): - async for _ in enc.encode( - [np.zeros((64, 64, 4), dtype=np.uint8)]): + async for _ in enc.encode([np.zeros((64, 64, 4), dtype=np.uint8)]): pass asyncio.run(run()) def test_encoder_close_is_idempotent(): + async def run(): - enc = FragmentedMP4Encoder( - width=64, height=64, fps=24, segment_idx=0) + enc = FragmentedMP4Encoder(width=64, height=64, fps=24, segment_idx=0) await enc.__aenter__() await enc.close() await enc.close() # no raise diff --git a/fastvideo/tests/entrypoints/streaming/test_worker.py b/fastvideo/tests/entrypoints/streaming/test_worker.py index 2a448636c..cb9ee2d8b 100644 --- a/fastvideo/tests/entrypoints/streaming/test_worker.py +++ b/fastvideo/tests/entrypoints/streaming/test_worker.py @@ -36,7 +36,10 @@ class _RecordingGenerator: state_to_return: ContinuationState | None = field(default_factory=lambda: ContinuationState( kind="ltx2.v1", - payload={"schema_version": 1, "segment_index": 1}, + payload={ + "schema_version": 1, + "segment_index": 1 + }, )) requests: list[GenerationRequest] = field(default_factory=list) diff --git a/fastvideo/tests/entrypoints/test_openai_api.py b/fastvideo/tests/entrypoints/test_openai_api.py index e4f7179a6..8097cac88 100644 --- a/fastvideo/tests/entrypoints/test_openai_api.py +++ b/fastvideo/tests/entrypoints/test_openai_api.py @@ -134,8 +134,7 @@ class TestImageBuildGenerationKwargs: def test_output_path_under_images_subdir(self, tmp_path): kw = self._build() - assert kw["output_path"].startswith( - os.path.join(str(tmp_path), "images")) + assert kw["output_path"].startswith(os.path.join(str(tmp_path), "images")) def test_num_frames_always_one(self): kw = self._build() @@ -189,8 +188,7 @@ class TestVideoBuildGenerationKwargs: def test_output_path_under_videos_subdir(self, tmp_path): kw = self._build() - assert kw["output_path"].startswith( - os.path.join(str(tmp_path), "videos")) + assert kw["output_path"].startswith(os.path.join(str(tmp_path), "videos")) assert kw["output_path"].endswith(".mp4") def test_fps_defaults_24(self): @@ -242,10 +240,8 @@ class TestVideoDefaultRequestMerge: body_defaults = dict(prompt="a running dog", seconds=4) body_defaults.update(body_overrides) req = VideoGenerationsRequest(**body_defaults) - default_request = _make_default_request( - default_raw) if default_raw else None - return _build_generation_kwargs( - "req-v1", req, default_request=default_request) + default_request = _make_default_request(default_raw) if default_raw else None + return _build_generation_kwargs("req-v1", req, default_request=default_request) def test_default_seed_flows_through_when_body_omits(self): kw = self._build(default_raw={"sampling": {"seed": 42}}) @@ -281,24 +277,17 @@ class TestVideoDefaultRequestMerge: def test_body_size_overrides_default_sampling_dims(self): kw = self._build( - default_raw={ - "sampling": { - "width": 640, - "height": 360 - } - }, + default_raw={"sampling": { + "width": 640, + "height": 360 + }}, size="1024x576", ) assert kw["width"] == 1024 assert kw["height"] == 576 def test_default_width_height_preserved_when_body_omits_size(self): - kw = self._build(default_raw={ - "sampling": { - "width": 640, - "height": 360 - } - }) + kw = self._build(default_raw={"sampling": {"width": 640, "height": 360}}) assert kw["width"] == 640 assert kw["height"] == 360 @@ -320,8 +309,7 @@ class TestVideoDefaultRequestMerge: assert default_dir not in kw["output_path"] def test_default_negative_prompt_flows_through(self): - kw = self._build( - default_raw={"negative_prompt": "low quality, blur"}) + kw = self._build(default_raw={"negative_prompt": "low quality, blur"}) assert kw["negative_prompt"] == "low quality, blur" def test_body_negative_prompt_overrides_default(self): @@ -339,16 +327,11 @@ class TestVideoDefaultRequestMerge: def test_default_request_not_mutated_by_build(self): # Merge should operate on a fresh copy (caller supplies a clone); # the helper itself must not mutate the passed-in default. - default_request = _make_default_request( - {"sampling": { - "seed": 42, - "fps": 30 - }}) + default_request = _make_default_request({"sampling": {"seed": 42, "fps": 30}}) from fastvideo.entrypoints.openai.video_api import ( _build_generation_kwargs, ) req = VideoGenerationsRequest(prompt="p", seconds=1) - _ = _build_generation_kwargs( - "req-1", req, default_request=default_request) + _ = _build_generation_kwargs("req-1", req, default_request=default_request) assert default_request.sampling.seed == 42 assert default_request.sampling.fps == 30 @@ -369,48 +352,30 @@ class TestValidateDefaultRequestAgainstPreset: def test_empty_stage_overrides_is_noop(self): from fastvideo.entrypoints.openai.api_server import ( _validate_default_request_against_preset, ) - default_request = _make_default_request({ - "sampling": { - "seed": 42 - } - }) + default_request = _make_default_request({"sampling": {"seed": 42}}) _validate_default_request_against_preset(default_request, "any/model") def test_unknown_model_path_is_noop(self): from fastvideo.entrypoints.openai.api_server import ( _validate_default_request_against_preset, ) - default_request = _make_default_request({ - "stage_overrides": { - "denoise": { - "num_inference_steps": 10 - } - } - }) + default_request = _make_default_request({"stage_overrides": {"denoise": {"num_inference_steps": 10}}}) with patch( "fastvideo.entrypoints.openai.api_server.get_preset_selection", return_value=(None, None), ): - _validate_default_request_against_preset( - default_request, "unknown/model") + _validate_default_request_against_preset(default_request, "unknown/model") def test_unknown_stage_name_raises(self): from fastvideo.api.errors import ConfigValidationError from fastvideo.entrypoints.openai.api_server import ( _validate_default_request_against_preset, ) - default_request = _make_default_request({ - "stage_overrides": { - "not_a_real_stage": { - "num_inference_steps": 10 - } - } - }) + default_request = _make_default_request({"stage_overrides": {"not_a_real_stage": {"num_inference_steps": 10}}}) with patch( "fastvideo.entrypoints.openai.api_server.get_preset_selection", return_value=("wan_t2v_1_3b", "wan"), ): with pytest.raises(ConfigValidationError): - _validate_default_request_against_preset( - default_request, "Wan-AI/Wan2.1-T2V-1.3B-Diffusers") + _validate_default_request_against_preset(default_request, "Wan-AI/Wan2.1-T2V-1.3B-Diffusers") # --------------------------------------------------------------------------- @@ -435,10 +400,7 @@ class TestDefaultRequestState: from fastvideo.entrypoints.openai import state as state_mod saved = state_mod._default_request try: - state_mod._default_request = _make_default_request( - {"sampling": { - "seed": 7 - }}) + state_mod._default_request = _make_default_request({"sampling": {"seed": 7}}) state_mod.clear_state() assert state_mod.get_default_request() is None finally: @@ -482,8 +444,6 @@ class TestProtocolModels: assert len(ids) == 100 def test_video_list_response(self): - resp = VideoListResponse( - data=[VideoResponse( - id="a"), VideoResponse(id="b")]) + resp = VideoListResponse(data=[VideoResponse(id="a"), VideoResponse(id="b")]) assert len(resp.data) == 2 assert resp.object == "list" diff --git a/fastvideo/tests/entrypoints/test_openai_api_integration.py b/fastvideo/tests/entrypoints/test_openai_api_integration.py index da35b4bc5..078434e85 100644 --- a/fastvideo/tests/entrypoints/test_openai_api_integration.py +++ b/fastvideo/tests/entrypoints/test_openai_api_integration.py @@ -97,21 +97,17 @@ def server(): pass # Check if the process died. if proc.poll() is not None: - stdout = proc.stdout.read().decode( - errors="replace") if proc.stdout else "" - pytest.fail( - f"Server exited with code {proc.returncode} " - f"before becoming healthy.\n--- stdout ---\n{stdout}") + stdout = proc.stdout.read().decode(errors="replace") if proc.stdout else "" + pytest.fail(f"Server exited with code {proc.returncode} " + f"before becoming healthy.\n--- stdout ---\n{stdout}") time.sleep(2) if not healthy: - stdout = proc.stdout.read().decode( - errors="replace") if proc.stdout else "" + stdout = proc.stdout.read().decode(errors="replace") if proc.stdout else "" proc.kill() proc.wait() - pytest.fail( - "Server did not become healthy within " - f"{SERVER_STARTUP_TIMEOUT_S}s.\n--- stdout ---\n{stdout}") + pytest.fail("Server did not become healthy within " + f"{SERVER_STARTUP_TIMEOUT_S}s.\n--- stdout ---\n{stdout}") yield {"base_url": base_url, "output_dir": output_dir} @@ -181,19 +177,16 @@ class TestVideoGeneration: status = job["status"] while status != "completed" and time.monotonic() < deadline: time.sleep(POLL_INTERVAL_S) - poll_resp = requests.get(f"{base_url}/v1/videos/{video_id}", - timeout=10) + poll_resp = requests.get(f"{base_url}/v1/videos/{video_id}", timeout=10) assert poll_resp.status_code == 200 status = poll_resp.json()["status"] if status == "failed": pytest.fail(f"Video generation failed: {poll_resp.json()}") - assert status == "completed", ( - f"Video not completed within {GENERATION_TIMEOUT_S}s") + assert status == "completed", (f"Video not completed within {GENERATION_TIMEOUT_S}s") # 3. Verify the video file exists on the server filesystem. - detail = requests.get(f"{base_url}/v1/videos/{video_id}", - timeout=10).json() + detail = requests.get(f"{base_url}/v1/videos/{video_id}", timeout=10).json() assert detail.get("file_path") is not None assert detail["file_path"].endswith(".mp4") diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index 6d64ced98..d123f60a3 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -250,7 +250,7 @@ def test_generate_single_video_metadata_only_skips_output_materialization(monkey assert result["frames"] is None assert result["video_path"] is None assert result["peak_memory_mb"] == 42.0 - assert empty_calls == [((0,), {"device": "cpu"})] + assert empty_calls == [((0, ), {"device": "cpu"})] def test_generate_single_video_return_frames_still_materializes_output(tmp_path): @@ -306,7 +306,7 @@ def test_generate_single_video_save_video_still_builds_frames(monkeypatch, tmp_p def test_generate_single_video_audio_only_metadata_returns_audio_without_frames(tmp_path): - audio = torch.zeros((16,), dtype=torch.float32) + audio = torch.zeros((16, ), dtype=torch.float32) output_batch = _single_video_output_batch( _NoCpuMaterializationOutput(), extra={ @@ -332,7 +332,7 @@ def test_generate_single_video_audio_only_metadata_returns_audio_without_frames( def test_generate_single_video_audio_only_save_skips_placeholder_materialization(tmp_path): - audio = torch.zeros((16,), dtype=torch.float32) + audio = torch.zeros((16, ), dtype=torch.float32) output_batch = _single_video_output_batch( _NoCpuMaterializationOutput(), extra={ @@ -360,7 +360,7 @@ def test_generate_single_video_audio_only_save_skips_placeholder_materialization def test_generate_single_video_ray_audio_only_save_preserves_worker_metadata(monkeypatch, tmp_path): - audio = torch.zeros((16,), dtype=torch.float32) + audio = torch.zeros((16, ), dtype=torch.float32) worker_output = ForwardBatch( data_type="audio", output=torch.ones((1, 3, 1, 8, 8)), diff --git a/fastvideo/tests/eval/conftest.py b/fastvideo/tests/eval/conftest.py index 8c5f0affc..365e50047 100644 --- a/fastvideo/tests/eval/conftest.py +++ b/fastvideo/tests/eval/conftest.py @@ -43,13 +43,13 @@ def _reference_metric_names() -> list[str]: # perfect-match value on duplicate (gen, ref). No reference JSON needed. INVARIANT_SPEC = [ # (metric, expected, abs_tolerance, location) - ("common.ssim", 1.0, 1e-3, "per_sample"), - ("common.psnr", 99.0, None, "per_sample"), # ≥ threshold (identical → ~inf, capped by floats) - ("common.lpips", 0.0, 1e-2, "per_sample"), + ("common.ssim", 1.0, 1e-3, "per_sample"), + ("common.psnr", 99.0, None, "per_sample"), # ≥ threshold (identical → ~inf, capped by floats) + ("common.lpips", 0.0, 1e-2, "per_sample"), ("optical_flow.gt_optical_flow", 0.0, 1e-2, "per_sample"), - ("common.fvd", 0.0, 1.0, "corpus"), - ("audio.frechet_distance", 0.0, 1.0, "corpus"), - ("audio.kl_divergence", 0.0, 1e-3, "per_sample"), + ("common.fvd", 0.0, 1.0, "corpus"), + ("audio.frechet_distance", 0.0, 1.0, "corpus"), + ("audio.kl_divergence", 0.0, 1e-3, "per_sample"), ] @@ -102,8 +102,10 @@ def invariant_results(asset_meta) -> EvalResults: """ tmp = Path(tempfile.mkdtemp(prefix="eval_regression_inv_")) try: - gen = tmp / "gen"; ref = tmp / "ref" - gen.mkdir(); ref.mkdir() + gen = tmp / "gen" + ref = tmp / "ref" + gen.mkdir() + ref.mkdir() for i in range(N_DUP): shutil.copyfile(ASSET_VIDEO, gen / f"clip_{i:03d}.mp4") shutil.copyfile(ASSET_VIDEO, ref / f"clip_{i:03d}.mp4") @@ -112,7 +114,8 @@ def invariant_results(asset_meta) -> EvalResults: skip_missing_deps=True, ) samples = samples_from( - video=gen, reference=ref, + video=gen, + reference=ref, text_prompt=asset_meta["text_prompt"], fps=asset_meta["fps"], extract_audio=tmp / "audio_cache", diff --git a/fastvideo/tests/eval/test_metric_identity_invariants.py b/fastvideo/tests/eval/test_metric_identity_invariants.py index ba62f7199..58623fe4d 100644 --- a/fastvideo/tests/eval/test_metric_identity_invariants.py +++ b/fastvideo/tests/eval/test_metric_identity_invariants.py @@ -15,8 +15,8 @@ REQUIRED_GPUS = 1 @pytest.mark.parametrize("metric,expected,tol,where", INVARIANT_SPEC) -def test_metric_identity_invariant(metric: str, expected: float, tol: float | None, - where: str, invariant_results) -> None: +def test_metric_identity_invariant(metric: str, expected: float, tol: float | None, where: str, + invariant_results) -> None: if where == "corpus": result = invariant_results.corpus.get(metric) else: @@ -30,11 +30,7 @@ def test_metric_identity_invariant(metric: str, expected: float, tol: float | No if tol is None: # threshold form (e.g. PSNR ≥ 99); for identical inputs PSNR # approaches infinity but float epsilon caps it at ~100. - assert result.score >= expected, ( - f"{metric}: {result.score:.4f} below threshold {expected}" - ) + assert result.score >= expected, (f"{metric}: {result.score:.4f} below threshold {expected}") else: delta = abs(result.score - expected) - assert delta <= tol, ( - f"{metric}: got {result.score:.6f}, expected ≈ {expected} (|Δ|={delta:.6f} > tol={tol})" - ) + assert delta <= tol, (f"{metric}: got {result.score:.6f}, expected ≈ {expected} (|Δ|={delta:.6f} > tol={tol})") diff --git a/fastvideo/tests/eval/test_metric_score_regression.py b/fastvideo/tests/eval/test_metric_score_regression.py index 239177355..c75eb6c93 100644 --- a/fastvideo/tests/eval/test_metric_score_regression.py +++ b/fastvideo/tests/eval/test_metric_score_regression.py @@ -16,29 +16,28 @@ from fastvideo.tests.eval.conftest import _reference_metric_names, load_referenc REQUIRED_GPUS = 1 - # Per-metric absolute tolerances. Tight for closed-form math, looser for # anything that runs CLIP/DINO/ViCLIP/VLM under bf16 or cuDNN. TOLERANCE: dict[str, float] = { # WER tolerates ~1 word/clip drift: reference runs openai-whisper with an # explicit language hint; fastvideo runs transformers WhisperForConditionalGeneration # without one (auto-detect). Same model weights, different decode paths. - "audio.wer": 0.2, - "audio.desync": 1e-3, - "audio.clap_score": 1e-2, - "audio.audiobox_aesthetics": 1e-2, - "audio.imagebind_score": 1e-2, + "audio.wer": 0.2, + "audio.desync": 1e-3, + "audio.clap_score": 1e-2, + "audio.audiobox_aesthetics": 1e-2, + "audio.imagebind_score": 1e-2, # videoscore2 is handled separately (per-dimension hard-score check) — # see _assert_videoscore2 below. - "vbench.aesthetic_quality": 1e-2, - "vbench.background_consistency": 1e-2, - "vbench.dynamic_degree": 0.0, # binary 0/1 — must match exactly - "vbench.imaging_quality": 1e-2, - "vbench.motion_smoothness": 1e-2, - "vbench.overall_consistency": 1e-2, - "vbench.subject_consistency": 1e-2, - "vbench.temporal_flickering": 1e-2, - "vbench.temporal_style": 1e-2, + "vbench.aesthetic_quality": 1e-2, + "vbench.background_consistency": 1e-2, + "vbench.dynamic_degree": 0.0, # binary 0/1 — must match exactly + "vbench.imaging_quality": 1e-2, + "vbench.motion_smoothness": 1e-2, + "vbench.overall_consistency": 1e-2, + "vbench.subject_consistency": 1e-2, + "vbench.temporal_flickering": 1e-2, + "vbench.temporal_style": 1e-2, } @@ -54,17 +53,14 @@ def _assert_videoscore2(result, reference: dict) -> None: which still fails loudly on a parse miss (hard score ``None``) or a ≥2 point divergence. """ - for dim in ("visual_quality_hard", "text_alignment_hard", - "physical_consistency_hard"): + for dim in ("visual_quality_hard", "text_alignment_hard", "physical_consistency_hard"): ref_h = reference["details"].get(dim) got_h = result.details.get(dim) assert ref_h is not None, f"reference missing {dim}" - assert got_h is not None, ( - f"videoscore2 did not parse {dim} " - f"(raw_output head: {result.details.get('raw_output', '')[:160]!r})") - assert abs(got_h - ref_h) <= 1, ( - f"videoscore2 {dim}: fastvideo={got_h} vs reference={ref_h} " - f"(integer sub-score drifted >1 across environments)") + assert got_h is not None, (f"videoscore2 did not parse {dim} " + f"(raw_output head: {result.details.get('raw_output', '')[:160]!r})") + assert abs(got_h - ref_h) <= 1, (f"videoscore2 {dim}: fastvideo={got_h} vs reference={ref_h} " + f"(integer sub-score drifted >1 across environments)") @pytest.mark.parametrize("metric_name", _reference_metric_names()) @@ -91,7 +87,5 @@ def test_metric_score_regression(metric_name: str, gold_results: dict) -> None: delta = abs(result.score - reference["score"]) assert math.isfinite(result.score), f"{metric_name} returned non-finite score" - assert delta <= tol, ( - f"{metric_name}: fastvideo={result.score:.6f} vs reference={reference['score']:.6f} " - f"(|Δ|={delta:.6f} > tol={tol})" - ) + assert delta <= tol, (f"{metric_name}: fastvideo={result.score:.6f} vs reference={reference['score']:.6f} " + f"(|Δ|={delta:.6f} > tol={tol})") diff --git a/fastvideo/tests/golden_gate/_harness.py b/fastvideo/tests/golden_gate/_harness.py index deb500024..0be65b4e0 100644 --- a/fastvideo/tests/golden_gate/_harness.py +++ b/fastvideo/tests/golden_gate/_harness.py @@ -40,8 +40,8 @@ os.environ.setdefault("LOCAL_RANK", "0") @dataclass class GateSpec: """Everything model-specific about one golden gate.""" - name: str # golden filename stem, e.g. "wan_t2v" - repo_id: str # HF checkpoint repo + name: str # golden filename stem, e.g. "wan_t2v" + repo_id: str # HF checkpoint repo build_block: Callable[[int], torch.nn.Module] make_inputs: Callable[[torch.device, int], dict[str, Any]] subfolder: str = "transformer" @@ -55,7 +55,7 @@ class GateSpec: dtype: torch.dtype = torch.bfloat16 # some blocks return tuples (e.g. double-stream img/txt) — reduce to one tensor postprocess: Callable[[Any], torch.Tensor] = field(default=lambda out: out) - model_root_env: str | None = None # env var naming a local checkpoint root + model_root_env: str | None = None # env var naming a local checkpoint root weight_file: str = "diffusion_pytorch_model.safetensors" # single-file fallback @@ -178,8 +178,13 @@ def run_gate(spec: GateSpec) -> None: out = out.detach().float().cpu() golden_path, golden_exists = _resolve_golden(spec) - meta = {"name": spec.name, "layer": spec.layer, "seed": spec.seed, - "shape": list(out.shape), "env": env_fingerprint(spec)} + meta = { + "name": spec.name, + "layer": spec.layer, + "seed": spec.seed, + "shape": list(out.shape), + "env": env_fingerprint(spec) + } if not golden_exists: golden_path.parent.mkdir(parents=True, exist_ok=True) torch.save({"output": out, "metadata": meta}, golden_path) @@ -192,16 +197,15 @@ def run_gate(spec: GateSpec) -> None: golden = torch.load(golden_path, weights_only=True) mismatches = { key: (value, meta["env"].get(key)) - for key, value in golden["metadata"]["env"].items() - if meta["env"].get(key) != value + for key, value in golden["metadata"]["env"].items() if meta["env"].get(key) != value } - assert not mismatches, ( - f"environment changed since golden was minted — re-mint deliberately rather " - f"than chasing bit drift. golden -> current: {mismatches}") + assert not mismatches, (f"environment changed since golden was minted — re-mint deliberately rather " + f"than chasing bit drift. golden -> current: {mismatches}") drift = (out - golden["output"]).abs() print(f"golden-gate[{spec.name}] drift: max_abs={drift.max().item():.10f} " - f"mean_abs={drift.mean().item():.10f}", flush=True) + f"mean_abs={drift.mean().item():.10f}", + flush=True) assert_close(out, golden["output"], atol=0.0, rtol=0.0) diff --git a/fastvideo/tests/golden_gate/test_dreamx.py b/fastvideo/tests/golden_gate/test_dreamx.py index 21238aa91..a61e21c90 100644 --- a/fastvideo/tests/golden_gate/test_dreamx.py +++ b/fastvideo/tests/golden_gate/test_dreamx.py @@ -29,8 +29,7 @@ TEXT_LEN = 512 def _build_block(layer: int) -> torch.nn.Module: - from fastvideo.configs.pipelines.dreamx_world import ( - make_dreamx_world_5b_cam_dit_config) + from fastvideo.configs.pipelines.dreamx_world import (make_dreamx_world_5b_cam_dit_config) from fastvideo.models.dits.dreamx_world import DreamXWorldTransformerBlock arch = make_dreamx_world_5b_cam_dit_config().arch_config @@ -38,18 +37,18 @@ def _build_block(layer: int) -> torch.nn.Module: # (dreamx_world.py:374-391) return DreamXWorldTransformerBlock( arch.num_attention_heads * arch.attention_head_dim, # 3072 - arch.ffn_dim, # 14336 - arch.num_attention_heads, # 24 - arch.qk_norm, # "rms_norm_across_heads" - arch.cross_attn_norm, # True - arch.eps, # 1e-6 - arch.added_kv_proj_dim, # None -> T2V cross-attn + arch.ffn_dim, # 14336 + arch.num_attention_heads, # 24 + arch.qk_norm, # "rms_norm_across_heads" + arch.cross_attn_norm, # True + arch.eps, # 1e-6 + arch.added_kv_proj_dim, # None -> T2V cross-attn arch._supported_attention_backends, quant_config=None, prefix=f"Wan.blocks.{layer}", - add_control_adapter=arch.add_control_adapter, # True - cam_method=arch.cam_method, # "prope" - attn_compress=arch.attn_compress, # 1 + add_control_adapter=arch.add_control_adapter, # True + cam_method=arch.cam_method, # "prope" + attn_compress=arch.attn_compress, # 1 cam_self_attn_layers=arch.cam_self_attn_layers, # None -> cam attn everywhere layer_idx=layer, ) @@ -59,21 +58,23 @@ def _make_inputs(device: torch.device, seed: int) -> dict: """Fixed batch mirroring the SSIM defaults (9 frames @ 64x64): latent (3, 4, 4) -> patch (1,2,2) -> seq = 3*2*2 = 12, cameras = 3 latent frames.""" from fastvideo.layers.rotary_embedding import get_rotary_pos_embed - from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import ( - build_dreamx_camera_condition) + from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import (build_dreamx_camera_condition) generator = torch.Generator(device="cpu").manual_seed(seed) num_frames, height, width = 9, 64, 64 - post_t = (num_frames - 1) // 4 + 1 # VAE 4x temporal + post_t = (num_frames - 1) // 4 + 1 # VAE 4x temporal post_h, post_w = height // 16 // 2, width // 16 // 2 # VAE 16x spatial, patch (1,2,2) - seq = post_t * post_h * post_w # 12 + seq = post_t * post_h * post_w # 12 # RoPE exactly as the model forward builds it (dreamx_world.py:429-439) d = INNER_DIM // 24 # head_dim 128 rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] # [44, 42, 42] - freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_t, post_h, post_w), INNER_DIM, 24, rope_dim_list, - rope_theta=10000, dtype=torch.float64) + freqs_cos, freqs_sin = get_rotary_pos_embed((post_t, post_h, post_w), + INNER_DIM, + 24, + rope_dim_list, + rope_theta=10000, + dtype=torch.float64) freqs_cis = (freqs_cos.to(device).float(), freqs_sin.to(device).float()) hidden = torch.randn(1, seq, INNER_DIM, generator=generator, dtype=torch.float32) @@ -83,9 +84,12 @@ def _make_inputs(device: torch.device, seed: int) -> dict: # Deterministic camera condition, same builder + keys as the pipeline stage # (stages.py:55-64). dtype MUST match hidden_states. - camera = build_dreamx_camera_condition( - ["w"], [4.0], num_frames=num_frames, height=height, width=width, - dtype=torch.bfloat16, device=device) + camera = build_dreamx_camera_condition(["w"], [4.0], + num_frames=num_frames, + height=height, + width=width, + dtype=torch.bfloat16, + device=device) y_camera = {k: v.unsqueeze(0) for k, v in camera.items()} # viewmats (1,3,4,4), K (1,3,3,3) return { diff --git a/fastvideo/tests/golden_gate/test_flux.py b/fastvideo/tests/golden_gate/test_flux.py index 99449b3e7..22babd762 100644 --- a/fastvideo/tests/golden_gate/test_flux.py +++ b/fastvideo/tests/golden_gate/test_flux.py @@ -30,8 +30,8 @@ def _build_block(layer: int) -> torch.nn.Module: arch = FluxTransformer2DArchConfig() return FluxTransformerBlock( dim=arch.num_attention_heads * arch.attention_head_dim, # 3072 - num_attention_heads=arch.num_attention_heads, # 24 - attention_head_dim=arch.attention_head_dim, # 128 + num_attention_heads=arch.num_attention_heads, # 24 + attention_head_dim=arch.attention_head_dim, # 128 supported_attention_backends=FluxTransformer2DModel._supported_attention_backends, ) diff --git a/fastvideo/tests/golden_gate/test_flux2_klein.py b/fastvideo/tests/golden_gate/test_flux2_klein.py index 01f70ef56..f223ed123 100644 --- a/fastvideo/tests/golden_gate/test_flux2_klein.py +++ b/fastvideo/tests/golden_gate/test_flux2_klein.py @@ -62,16 +62,12 @@ def _make_inputs(device: torch.device, seed: int) -> dict: # (shift, scale, gate) per set, [1, 1, DIM] each, in the DiT dtype. return tuple( tuple( - torch.randn(1, 1, DIM, generator=generator, dtype=torch.float32) - .to(device=device, dtype=torch.bfloat16) - for _ in range(3)) - for _ in range(n_sets)) + torch.randn(1, 1, DIM, generator=generator, dtype=torch.float32).to(device=device, dtype=torch.bfloat16) + for _ in range(3)) for _ in range(n_sets)) # ids built as in Flux2Transformer2DModel.forward (flux_2.py:1044-1051) - txt_ids = torch.cartesian_prod( - torch.arange(1), torch.arange(1), torch.arange(1), torch.arange(TEXT_LEN)) - img_ids = torch.cartesian_prod( - torch.arange(1), torch.arange(IMG_H), torch.arange(IMG_W), torch.arange(1)) + txt_ids = torch.cartesian_prod(torch.arange(1), torch.arange(1), torch.arange(1), torch.arange(TEXT_LEN)) + img_ids = torch.cartesian_prod(torch.arange(1), torch.arange(IMG_H), torch.arange(IMG_W), torch.arange(1)) rope = Flux2PosEmbed(theta=2000, axes_dim=[32, 32, 32, 32]) freqs_cis = compute_flux2_freqs_cis_from_ids(rope, txt_ids, img_ids, device=device) diff --git a/fastvideo/tests/golden_gate/test_gamecraft.py b/fastvideo/tests/golden_gate/test_gamecraft.py index 2bf8a42fc..6457071b7 100644 --- a/fastvideo/tests/golden_gate/test_gamecraft.py +++ b/fastvideo/tests/golden_gate/test_gamecraft.py @@ -30,10 +30,10 @@ def _build_block(layer: int) -> torch.nn.Module: arch = HunyuanGameCraftArchConfig() return MMDoubleStreamBlock( - hidden_size=arch.hidden_size, # 3072 - num_attention_heads=arch.num_attention_heads, # 24 - mlp_ratio=arch.mlp_ratio, # 4.0 - dtype=arch.dtype, # None + hidden_size=arch.hidden_size, # 3072 + num_attention_heads=arch.num_attention_heads, # 24 + mlp_ratio=arch.mlp_ratio, # 4.0 + dtype=arch.dtype, # None supported_attention_backends=arch._supported_attention_backends, prefix=f"hunyuan_gamecraft.double_blocks.{layer}", ) @@ -52,8 +52,7 @@ def _make_inputs(device: torch.device, seed: int) -> dict: tt, th, tw = 3, 4, 5 img_seq, txt_seq = tt * th * tw, 32 - freqs_cos, freqs_sin = get_rotary_pos_embed( - (tt, th, tw), HIDDEN, 24, [16, 56, 56], 256) # -> [60, 128] fp32 each + freqs_cos, freqs_sin = get_rotary_pos_embed((tt, th, tw), HIDDEN, 24, [16, 56, 56], 256) # -> [60, 128] fp32 each img = torch.randn(1, img_seq, HIDDEN, generator=generator, dtype=torch.float32) txt = torch.randn(1, txt_seq, HIDDEN, generator=generator, dtype=torch.float32) @@ -63,8 +62,8 @@ def _make_inputs(device: torch.device, seed: int) -> dict: "img": img.to(device=device, dtype=torch.bfloat16), "txt": txt.to(device=device, dtype=torch.bfloat16), "vec": vec.to(device=device, dtype=torch.bfloat16), - "freqs_cis": (freqs_cos.to(device=device, dtype=torch.bfloat16), - freqs_sin.to(device=device, dtype=torch.bfloat16)), + "freqs_cis": (freqs_cos.to(device=device, + dtype=torch.bfloat16), freqs_sin.to(device=device, dtype=torch.bfloat16)), } diff --git a/fastvideo/tests/golden_gate/test_gen3c.py b/fastvideo/tests/golden_gate/test_gen3c.py index 498426f2b..781e7312a 100644 --- a/fastvideo/tests/golden_gate/test_gen3c.py +++ b/fastvideo/tests/golden_gate/test_gen3c.py @@ -30,12 +30,12 @@ def _build_block(layer: int) -> torch.nn.Module: arch = Gen3CArchConfig() return Gen3CTransformerBlock( num_attention_heads=arch.num_attention_heads, # 32 - attention_head_dim=arch.attention_head_dim, # 128 -> hidden 4096 - cross_attention_dim=arch.text_embed_dim, # 1024 - mlp_ratio=arch.mlp_ratio, # 4.0 -> mlp inner 16384 - adaln_lora_dim=arch.adaln_lora_dim, # 256 - use_adaln_lora=arch.use_adaln_lora, # True - qk_norm=(arch.qk_norm == "rms_norm"), # True + attention_head_dim=arch.attention_head_dim, # 128 -> hidden 4096 + cross_attention_dim=arch.text_embed_dim, # 1024 + mlp_ratio=arch.mlp_ratio, # 4.0 -> mlp inner 16384 + adaln_lora_dim=arch.adaln_lora_dim, # 256 + use_adaln_lora=arch.use_adaln_lora, # True + qk_norm=(arch.qk_norm == "rms_norm"), # True supported_attention_backends=arch._supported_attention_backends, ) diff --git a/fastvideo/tests/golden_gate/test_glm_image.py b/fastvideo/tests/golden_gate/test_glm_image.py index 92ad32032..48b567dc5 100644 --- a/fastvideo/tests/golden_gate/test_glm_image.py +++ b/fastvideo/tests/golden_gate/test_glm_image.py @@ -28,15 +28,14 @@ LATENT_H = LATENT_W = 64 # patch_size=2 -> 32x32 = 1024 image tokens def _build_block(layer: int) -> torch.nn.Module: from fastvideo.configs.models.dits.glm_image import GlmImageDiTArchConfig - from fastvideo.models.dits.glm_image import (GlmImageTransformer2DModel, - GlmImageTransformerBlock) + from fastvideo.models.dits.glm_image import (GlmImageTransformer2DModel, GlmImageTransformerBlock) arch = GlmImageDiTArchConfig() return GlmImageTransformerBlock( arch.num_attention_heads * arch.attention_head_dim, # dim = 4096 - arch.num_attention_heads, # 32 - arch.attention_head_dim, # 128 - arch.time_embed_dim, # 512 + arch.num_attention_heads, # 32 + arch.attention_head_dim, # 128 + arch.time_embed_dim, # 512 supported_attention_backends=GlmImageTransformer2DModel._supported_attention_backends, prefix=f"transformer_blocks.{layer}", ) diff --git a/fastvideo/tests/golden_gate/test_kandinsky5.py b/fastvideo/tests/golden_gate/test_kandinsky5.py index 5f68c27d7..32102b8d1 100644 --- a/fastvideo/tests/golden_gate/test_kandinsky5.py +++ b/fastvideo/tests/golden_gate/test_kandinsky5.py @@ -73,8 +73,7 @@ def _make_inputs(device: torch.device, seed: int) -> dict: # matrices [B, T, H, W, 1, head_dim/2, 2, 2] fp32, then the dense # fractal_flatten(block_mask=False) is rope.flatten(1, 3) (line 109). rope3d = Kandinsky5RoPE3D(AXES_DIMS) - rope = rope3d((b, t, h, w), - [torch.arange(t), torch.arange(h), torch.arange(w)]) + rope = rope3d((b, t, h, w), [torch.arange(t), torch.arange(h), torch.arange(w)]) rope = rope.flatten(1, 3) return { diff --git a/fastvideo/tests/golden_gate/test_lingbot.py b/fastvideo/tests/golden_gate/test_lingbot.py index 6d5cda55b..5e28788a7 100644 --- a/fastvideo/tests/golden_gate/test_lingbot.py +++ b/fastvideo/tests/golden_gate/test_lingbot.py @@ -53,8 +53,12 @@ def _make_inputs(device: torch.device, seed: int) -> dict: d = INNER_DIM // 40 rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] - freqs_cos, freqs_sin = get_rotary_pos_embed( - (t, h, w), INNER_DIM, 40, rope_dim_list, rope_theta=10000, dtype=torch.float64) + freqs_cos, freqs_sin = get_rotary_pos_embed((t, h, w), + INNER_DIM, + 40, + rope_dim_list, + rope_theta=10000, + dtype=torch.float64) freqs_cis = (freqs_cos.to(device).float(), freqs_sin.to(device).float()) hidden = torch.randn(1, seq, INNER_DIM, generator=generator, dtype=torch.float32) diff --git a/fastvideo/tests/golden_gate/test_longcat.py b/fastvideo/tests/golden_gate/test_longcat.py index 1da7486e0..5ea910bde 100644 --- a/fastvideo/tests/golden_gate/test_longcat.py +++ b/fastvideo/tests/golden_gate/test_longcat.py @@ -23,7 +23,7 @@ from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, __all__ = ["distributed_runtime"] HIDDEN_SIZE = 4096 # configs/models/dits/longcat.py:77 -ADALN_DIM = 512 # configs/models/dits/longcat.py:93 +ADALN_DIM = 512 # configs/models/dits/longcat.py:93 def _build_block(layer: int) -> torch.nn.Module: diff --git a/fastvideo/tests/golden_gate/test_ltx2.py b/fastvideo/tests/golden_gate/test_ltx2.py index 32da8aca8..c2834481d 100644 --- a/fastvideo/tests/golden_gate/test_ltx2.py +++ b/fastvideo/tests/golden_gate/test_ltx2.py @@ -34,8 +34,7 @@ AUDIO_DIM = 2048 # 32 heads x 64 def _build_block(layer: int) -> torch.nn.Module: from fastvideo.configs.models.dits.ltx2 import LTX2VideoArchConfig - from fastvideo.models.dits.ltx2 import ( - BasicAVTransformerBlock, LTXRopeType, TransformerConfig) + from fastvideo.models.dits.ltx2 import (BasicAVTransformerBlock, LTXRopeType, TransformerConfig) arch = LTX2VideoArchConfig() video_cfg = TransformerConfig( @@ -69,13 +68,11 @@ def _build_block(layer: int) -> torch.nn.Module: def _make_inputs(device: torch.device, seed: int) -> dict: - from fastvideo.models.dits.ltx2 import ( - AudioLatentPatchifier, AudioLatentShape, DEFAULT_LTX2_AUDIO_DOWNSAMPLE, - DEFAULT_LTX2_AUDIO_HOP_LENGTH, DEFAULT_LTX2_AUDIO_MEL_BINS, - DEFAULT_LTX2_AUDIO_SAMPLE_RATE, DEFAULT_LTX2_SCALE_FACTORS, LTXRopeType, - TransformerArgs, VideoLatentPatchifier, VideoLatentShape, - _get_pixel_coords, generate_ltx_freq_grid_float64, - precompute_ltx_freqs_cis) + from fastvideo.models.dits.ltx2 import (AudioLatentPatchifier, AudioLatentShape, DEFAULT_LTX2_AUDIO_DOWNSAMPLE, + DEFAULT_LTX2_AUDIO_HOP_LENGTH, DEFAULT_LTX2_AUDIO_MEL_BINS, + DEFAULT_LTX2_AUDIO_SAMPLE_RATE, DEFAULT_LTX2_SCALE_FACTORS, LTXRopeType, + TransformerArgs, VideoLatentPatchifier, VideoLatentShape, _get_pixel_coords, + generate_ltx_freq_grid_float64, precompute_ltx_freqs_cis) generator = torch.Generator(device="cpu").manual_seed(seed) B, frames, h, w = 1, 2, 4, 6 @@ -84,67 +81,60 @@ def _make_inputs(device: torch.device, seed: int) -> dict: # Video positions: patch grid bounds -> pixel coords (scale 8,32,32; # fps 24; causal fix), cast to hidden dtype like ltx2.py:3044. - v_pos = VideoLatentPatchifier(patch_size=1).get_patch_grid_bounds( - VideoLatentShape((B, 128, frames, h, w)), device=torch.device("cpu")) - v_pos = _get_pixel_coords(v_pos, DEFAULT_LTX2_SCALE_FACTORS, fps=24.0, - causal_fix=True).to(torch.bfloat16) + v_pos = VideoLatentPatchifier(patch_size=1).get_patch_grid_bounds(VideoLatentShape((B, 128, frames, h, w)), + device=torch.device("cpu")) + v_pos = _get_pixel_coords(v_pos, DEFAULT_LTX2_SCALE_FACTORS, fps=24.0, causal_fix=True).to(torch.bfloat16) # Audio positions stay float32 [B, 1, T, 2] (ltx2.py:3099-3101). a_pos = AudioLatentPatchifier( patch_size=DEFAULT_LTX2_AUDIO_MEL_BINS, sample_rate=DEFAULT_LTX2_AUDIO_SAMPLE_RATE, hop_length=DEFAULT_LTX2_AUDIO_HOP_LENGTH, audio_latent_downsample_factor=DEFAULT_LTX2_AUDIO_DOWNSAMPLE, - is_causal=True, shift=0, - ).get_patch_grid_bounds( - AudioLatentShape((B, 8, seq_a, DEFAULT_LTX2_AUDIO_MEL_BINS)), - device=torch.device("cpu")) + is_causal=True, + shift=0, + ).get_patch_grid_bounds(AudioLatentShape((B, 8, seq_a, DEFAULT_LTX2_AUDIO_MEL_BINS)), device=torch.device("cpu")) - rope_kw = dict(out_dtype=torch.bfloat16, theta=10000.0, - use_middle_indices_grid=True, num_attention_heads=32, + rope_kw = dict(out_dtype=torch.bfloat16, + theta=10000.0, + use_middle_indices_grid=True, + num_attention_heads=32, rope_type=LTXRopeType.SPLIT, freq_grid_generator=generate_ltx_freq_grid_float64) - v_pe = precompute_ltx_freqs_cis(v_pos, dim=VIDEO_DIM, - max_pos=[20, 2048, 2048], **rope_kw) - v_cross_pe = precompute_ltx_freqs_cis(v_pos[:, 0:1, :], dim=AUDIO_DIM, - max_pos=[20], **rope_kw) - a_pe = precompute_ltx_freqs_cis(a_pos, dim=AUDIO_DIM, max_pos=[20], - **rope_kw) - a_cross_pe = precompute_ltx_freqs_cis(a_pos[:, 0:1, :], dim=AUDIO_DIM, - max_pos=[20], **rope_kw) + v_pe = precompute_ltx_freqs_cis(v_pos, dim=VIDEO_DIM, max_pos=[20, 2048, 2048], **rope_kw) + v_cross_pe = precompute_ltx_freqs_cis(v_pos[:, 0:1, :], dim=AUDIO_DIM, max_pos=[20], **rope_kw) + a_pe = precompute_ltx_freqs_cis(a_pos, dim=AUDIO_DIM, max_pos=[20], **rope_kw) + a_cross_pe = precompute_ltx_freqs_cis(a_pos[:, 0:1, :], dim=AUDIO_DIM, max_pos=[20], **rope_kw) def rnd(*shape): - return torch.randn(*shape, generator=generator, - dtype=torch.float32).to(device=device, - dtype=torch.bfloat16) + return torch.randn(*shape, generator=generator, dtype=torch.float32).to(device=device, dtype=torch.bfloat16) def to_dev(pe): return tuple(t.to(device) for t in pe) - video = TransformerArgs( - x=rnd(B, seq_v, VIDEO_DIM), - context=rnd(B, n_text, VIDEO_DIM), - context_mask=None, - timesteps=rnd(B, seq_v, 6 * VIDEO_DIM), - embedded_timestep=rnd(B, seq_v, VIDEO_DIM), - positional_embeddings=to_dev(v_pe), - cross_positional_embeddings=to_dev(v_cross_pe), - cross_scale_shift_timestep=rnd(B, seq_v, 4 * VIDEO_DIM), - cross_gate_timestep=rnd(B, seq_v, VIDEO_DIM), - enabled=True, prompt_timestep=None) - audio = TransformerArgs( - x=rnd(B, seq_a, AUDIO_DIM), - context=rnd(B, n_text, AUDIO_DIM), - context_mask=None, - timesteps=rnd(B, seq_a, 6 * AUDIO_DIM), - embedded_timestep=rnd(B, seq_a, AUDIO_DIM), - positional_embeddings=to_dev(a_pe), - cross_positional_embeddings=to_dev(a_cross_pe), - cross_scale_shift_timestep=rnd(B, seq_a, 4 * AUDIO_DIM), - cross_gate_timestep=rnd(B, seq_a, AUDIO_DIM), - enabled=True, prompt_timestep=None) + video = TransformerArgs(x=rnd(B, seq_v, VIDEO_DIM), + context=rnd(B, n_text, VIDEO_DIM), + context_mask=None, + timesteps=rnd(B, seq_v, 6 * VIDEO_DIM), + embedded_timestep=rnd(B, seq_v, VIDEO_DIM), + positional_embeddings=to_dev(v_pe), + cross_positional_embeddings=to_dev(v_cross_pe), + cross_scale_shift_timestep=rnd(B, seq_v, 4 * VIDEO_DIM), + cross_gate_timestep=rnd(B, seq_v, VIDEO_DIM), + enabled=True, + prompt_timestep=None) + audio = TransformerArgs(x=rnd(B, seq_a, AUDIO_DIM), + context=rnd(B, n_text, AUDIO_DIM), + context_mask=None, + timesteps=rnd(B, seq_a, 6 * AUDIO_DIM), + embedded_timestep=rnd(B, seq_a, AUDIO_DIM), + positional_embeddings=to_dev(a_pe), + cross_positional_embeddings=to_dev(a_cross_pe), + cross_scale_shift_timestep=rnd(B, seq_a, 4 * AUDIO_DIM), + cross_gate_timestep=rnd(B, seq_a, AUDIO_DIM), + enabled=True, + prompt_timestep=None) - return {"video": video, "audio": audio, - "video_original_seq_len": seq_v, "audio_original_seq_len": seq_a} + return {"video": video, "audio": audio, "video_original_seq_len": seq_v, "audio_original_seq_len": seq_a} SPEC = GateSpec( diff --git a/fastvideo/tests/golden_gate/test_matrixgame.py b/fastvideo/tests/golden_gate/test_matrixgame.py index 0832b1f5b..ab65badcd 100644 --- a/fastvideo/tests/golden_gate/test_matrixgame.py +++ b/fastvideo/tests/golden_gate/test_matrixgame.py @@ -42,7 +42,6 @@ MG3_INNER_DIM = 3072 # 24 heads x 128, matrixgame3.py:42-43 F, H, W = 3, 22, 40 # post-patch grid; seq=2640 (% 128 != 0, see docstring) N_PIX = 4 * (F - 1) + 1 # pixel-frame action length (action_module.py:511-521) - # ===== Matrix-Game-2.0 (causal) ===== @@ -63,14 +62,14 @@ def _build_block_mg2(layer: int) -> torch.nn.Module: assert inner_dim == MG2_INNER_DIM return CausalMatrixGame2TransformerBlock( inner_dim, - hf["ffn_dim"], # 8960 - hf["num_attention_heads"], # 12 - hf["local_attn_size"], # 6 - hf["sink_size"], # 0 - hf["qk_norm"], # "rms_norm_across_heads" - arch.cross_attn_norm, # True (block asserts it, causal_model.py:436) - hf["eps"], # 1e-6 - arch.added_kv_proj_dim, # None -> attn2 = WanT2VCrossAttention + hf["ffn_dim"], # 8960 + hf["num_attention_heads"], # 12 + hf["local_attn_size"], # 6 + hf["sink_size"], # 0 + hf["qk_norm"], # "rms_norm_across_heads" + arch.cross_attn_norm, # True (block asserts it, causal_model.py:436) + hf["eps"], # 1e-6 + arch.added_kv_proj_dim, # None -> attn2 = WanT2VCrossAttention arch._supported_attention_backends, prefix=f"Wan.blocks.{layer}", action_config=hf["action_config"], # ActionModule on blocks 0-14 only @@ -89,16 +88,22 @@ def _make_inputs_mg2(device: torch.device, seed: int) -> dict: # CLIP image embeds post WanImageEmbedding context = torch.randn(1, 257, MG2_INNER_DIM, generator=generator, dtype=torch.float32) temb = torch.randn(1, F, 6, MG2_INNER_DIM, generator=generator, dtype=torch.float32) - mouse = torch.randn(1, N_PIX, 2, generator=generator, dtype=torch.float32) # mouse_dim_in=2 + mouse = torch.randn(1, N_PIX, 2, generator=generator, dtype=torch.float32) # mouse_dim_in=2 keyboard = torch.randn(1, N_PIX, 4, generator=generator, dtype=torch.float32) # keyboard_dim_in=4 # masks exactly as _forward_inference builds them (causal_model.py:1129-1162) - block_mask = CausalMatrixGame2WanModel._prepare_blockwise_causal_attn_mask( - device=device, num_frames=F, frame_seqlen=H * W, - num_frame_per_block=3, local_attn_size=6, sink_size=0) - block_mask_action = CausalMatrixGame2WanModel._prepare_blockwise_causal_attn_mask_action( - device=device, num_frames=F, frame_seqlen=1, - num_frame_per_block=3, local_attn_size=6, sink_size=0) + block_mask = CausalMatrixGame2WanModel._prepare_blockwise_causal_attn_mask(device=device, + num_frames=F, + frame_seqlen=H * W, + num_frame_per_block=3, + local_attn_size=6, + sink_size=0) + block_mask_action = CausalMatrixGame2WanModel._prepare_blockwise_causal_attn_mask_action(device=device, + num_frames=F, + frame_seqlen=1, + num_frame_per_block=3, + local_attn_size=6, + sink_size=0) return { "hidden_states": hidden.to(device=device, dtype=torch.float32), @@ -145,16 +150,16 @@ def _build_block_mg3(layer: int) -> torch.nn.Module: assert inner_dim == MG3_INNER_DIM return MatrixGame3TransformerBlock( inner_dim, - arch.ffn_dim, # 14336 - arch.num_attention_heads, # 24 - arch.qk_norm, # "rms_norm_across_heads" - arch.cross_attn_norm, # True (asserted at model.py:246) - arch.eps, # 1e-6 + arch.ffn_dim, # 14336 + arch.num_attention_heads, # 24 + arch.qk_norm, # "rms_norm_across_heads" + arch.cross_attn_norm, # True (asserted at model.py:246) + arch.eps, # 1e-6 arch._supported_attention_backends, prefix=f"Wan.blocks.{layer}", action_config=arch.action_config, # blocks 0-14 only block_id=layer, - use_memory=arch.use_memory, # True -> cam_* Linears on EVERY block + use_memory=arch.use_memory, # True -> cam_* Linears on EVERY block ) @@ -169,12 +174,12 @@ def _make_inputs_mg3(device: torch.device, seed: int) -> dict: context = torch.randn(1, 512, MG3_INNER_DIM, generator=generator, dtype=torch.float32) # per-token 4-D temb (bs, seq, 6, 3072), real path (model.py:687-713) temb = torch.randn(1, seq, 6, MG3_INNER_DIM, generator=generator, dtype=torch.float32) - mouse = torch.randn(1, N_PIX, 2, generator=generator, dtype=torch.float32) # mouse_dim_in=2 + mouse = torch.randn(1, N_PIX, 2, generator=generator, dtype=torch.float32) # mouse_dim_in=2 keyboard = torch.randn(1, N_PIX, 6, generator=generator, dtype=torch.float32) # keyboard_dim_in=6 # exactly as model.py:640-646: use_memory=True -> per-head 3-D complex table - freqs = _build_rope_freqs(max_seq_len=2048, head_dim=128, num_heads=24, - sigma_theta=0.8, device=device) # (24, 2048, 64) complex128 + freqs = _build_rope_freqs(max_seq_len=2048, head_dim=128, num_heads=24, sigma_theta=0.8, + device=device) # (24, 2048, 64) complex128 return { "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), diff --git a/fastvideo/tests/golden_gate/test_minimax_h3_t2v.py b/fastvideo/tests/golden_gate/test_minimax_h3_t2v.py index 00b6c828a..675b5a508 100644 --- a/fastvideo/tests/golden_gate/test_minimax_h3_t2v.py +++ b/fastvideo/tests/golden_gate/test_minimax_h3_t2v.py @@ -110,8 +110,7 @@ def _component_dir() -> Path: from huggingface_hub import hf_hub_download - index_path = Path(hf_hub_download( - MODEL_REPO_ID, "transformer/diffusion_pytorch_model.safetensors.index.json")) + index_path = Path(hf_hub_download(MODEL_REPO_ID, "transformer/diffusion_pytorch_model.safetensors.index.json")) weight_map = json.loads(index_path.read_text())["weight_map"] prefix = f"transformer_blocks.{LAYER}." shards = sorted({shard for name, shard in weight_map.items() if name.startswith(prefix)}) @@ -173,9 +172,9 @@ def _make_inputs(device: torch.device) -> dict: seq = n_text + n_audio + n_video token_tags = torch.cat([ - torch.full((n_text,), 1, dtype=torch.long), - torch.full((n_audio,), 2, dtype=torch.long), - torch.full((n_video,), 0, dtype=torch.long), + torch.full((n_text, ), 1, dtype=torch.long), + torch.full((n_audio, ), 2, dtype=torch.long), + torch.full((n_video, ), 0, dtype=torch.long), ]) timestep_indices = torch.zeros(seq, dtype=torch.long) timestep_indices[n_text:n_text + n_audio] = 1 @@ -214,8 +213,7 @@ def _golden_filename() -> str: def _resolve_golden() -> tuple[Path, bool]: """Return (path, exists). Local dir wins; falls back to the HF dataset.""" - local_root = Path(os.environ.get( - "MINIMAX_H3_GATE_GOLDEN_DIR", Path(__file__).resolve().parent / "goldens")) + local_root = Path(os.environ.get("MINIMAX_H3_GATE_GOLDEN_DIR", Path(__file__).resolve().parent / "goldens")) local = local_root / _device_slug() / _golden_filename() if local.exists(): return local, True @@ -279,12 +277,10 @@ def test_minimax_h3_t2v_golden_gate() -> None: golden = torch.load(golden_path, weights_only=True) mismatches = { key: (value, meta["env"].get(key)) - for key, value in golden["metadata"]["env"].items() - if meta["env"].get(key) != value + for key, value in golden["metadata"]["env"].items() if meta["env"].get(key) != value } - assert not mismatches, ( - f"environment changed since golden was minted — re-mint deliberately rather " - f"than chasing bit drift. golden -> current: {mismatches}") + assert not mismatches, (f"environment changed since golden was minted — re-mint deliberately rather " + f"than chasing bit drift. golden -> current: {mismatches}") drift = (out - golden["output"]).abs() print(f"golden-gate drift: max_abs={drift.max().item():.10f} " diff --git a/fastvideo/tests/golden_gate/test_sd35.py b/fastvideo/tests/golden_gate/test_sd35.py index 659f36361..aff6e377c 100644 --- a/fastvideo/tests/golden_gate/test_sd35.py +++ b/fastvideo/tests/golden_gate/test_sd35.py @@ -30,8 +30,7 @@ N_TXT = 154 def _build_block(layer: int) -> torch.nn.Module: # Mirrors SD3Transformer2DModel.__init__ (sd3.py:923-933). from fastvideo.configs.models.dits.sd3 import SD3Transformer2DArchConfig - from fastvideo.models.dits.sd3 import (SD3JointTransformerBlock, - SD3Transformer2DModel) + from fastvideo.models.dits.sd3 import (SD3JointTransformerBlock, SD3Transformer2DModel) arch = SD3Transformer2DArchConfig() return SD3JointTransformerBlock( diff --git a/fastvideo/tests/golden_gate/test_stable_audio.py b/fastvideo/tests/golden_gate/test_stable_audio.py index 6da62edb7..af05c7910 100644 --- a/fastvideo/tests/golden_gate/test_stable_audio.py +++ b/fastvideo/tests/golden_gate/test_stable_audio.py @@ -23,7 +23,7 @@ from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, __all__ = ["distributed_runtime"] EMBED_DIM = 1536 # StableAudioArchConfig.embed_dim -COND_DIM = 768 # StableAudioArchConfig.cond_token_dim (project_cond_tokens=False) +COND_DIM = 768 # StableAudioArchConfig.cond_token_dim (project_cond_tokens=False) def _build_block(layer: int) -> torch.nn.Module: diff --git a/fastvideo/tests/golden_gate/test_wan_t2v.py b/fastvideo/tests/golden_gate/test_wan_t2v.py index e811e1901..8e8b99bd0 100644 --- a/fastvideo/tests/golden_gate/test_wan_t2v.py +++ b/fastvideo/tests/golden_gate/test_wan_t2v.py @@ -54,8 +54,12 @@ def _make_inputs(device: torch.device, seed: int) -> dict: d = INNER_DIM // 12 rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] - freqs_cos, freqs_sin = get_rotary_pos_embed( - (t, h, w), INNER_DIM, 12, rope_dim_list, rope_theta=10000, dtype=torch.float64) + freqs_cos, freqs_sin = get_rotary_pos_embed((t, h, w), + INNER_DIM, + 12, + rope_dim_list, + rope_theta=10000, + dtype=torch.float64) freqs_cis = (freqs_cos.to(device).float(), freqs_sin.to(device).float()) hidden = torch.randn(1, seq, INNER_DIM, generator=generator, dtype=torch.float32) diff --git a/fastvideo/tests/golden_gate/test_zimage.py b/fastvideo/tests/golden_gate/test_zimage.py index d282744d6..cd27b6774 100644 --- a/fastvideo/tests/golden_gate/test_zimage.py +++ b/fastvideo/tests/golden_gate/test_zimage.py @@ -52,8 +52,7 @@ def _make_inputs(device: torch.device, seed: int) -> dict: # freqs_cis exactly as the model builds it (zimage.py:342, 527-531); image # position ids start at frame index cap_len+1 -> use start=(1, 0, 0). rope = RopeEmbedder(theta=256.0, axes_dims=(32, 48, 48), axes_lens=(1536, 512, 512)) - pos_ids = ZImageTransformer2DModel.create_coordinate_grid( - (1, 8, 8), start=(1, 0, 0)).flatten(0, 2) # [64, 3] int32 + pos_ids = ZImageTransformer2DModel.create_coordinate_grid((1, 8, 8), start=(1, 0, 0)).flatten(0, 2) # [64, 3] int32 freqs_cis = rope(pos_ids).unsqueeze(0).expand(batch, seq, 64).contiguous() # complex64 [2, 64, 64] attn_mask = torch.zeros(batch, seq, dtype=torch.bool) diff --git a/fastvideo/tests/hooks/test_hooks_api.py b/fastvideo/tests/hooks/test_hooks_api.py index 5142f8449..4f6b277cf 100644 --- a/fastvideo/tests/hooks/test_hooks_api.py +++ b/fastvideo/tests/hooks/test_hooks_api.py @@ -5,6 +5,7 @@ import torch class EventHook(ForwardHook): + def __init__(self, content: str, event_list: list[str]): self.content = content self.event_list = event_list @@ -13,16 +14,12 @@ class EventHook(ForwardHook): return f"EventHook_{self.content}" def pre_forward(self, module: nn.Module, *args, **kwargs): - print( - f"[{self.content}] Pre-forward called with args[0].shape: {args[0].shape}" - ) + print(f"[{self.content}] Pre-forward called with args[0].shape: {args[0].shape}") self.event_list.append(f"[pre]{self.content}") return args, kwargs def post_forward(self, module: nn.Module, output: Any): - print( - f"[{self.content}] Post-forward called with outputs.shape: {output.shape}" - ) + print(f"[{self.content}] Post-forward called with outputs.shape: {output.shape}") self.event_list.append(f"[post]{self.content}") return output @@ -54,7 +51,5 @@ def test_hook_execution_order(): # Post-hooks should be LIFO (Last In First Out): B then A expected_events = ["[pre]A", "[pre]B", "[post]B", "[post]A"] - assert events == expected_events, ( - f"Expected {expected_events}, but got {events}" - ) + assert events == expected_events, (f"Expected {expected_events}, but got {events}") print(f"✓ Hook execution order test passed: {events}") diff --git a/fastvideo/tests/hooks/test_layerwise_offload.py b/fastvideo/tests/hooks/test_layerwise_offload.py index 3ad5d3144..f1505029e 100644 --- a/fastvideo/tests/hooks/test_layerwise_offload.py +++ b/fastvideo/tests/hooks/test_layerwise_offload.py @@ -36,9 +36,7 @@ class SimpleModelWithModuleList(nn.Module): dtype: torch.dtype = torch.float32, ): super().__init__() - self.blocks = nn.ModuleList( - [SimpleBlock(hidden_size, dtype=dtype) for _ in range(num_blocks)] - ) + self.blocks = nn.ModuleList([SimpleBlock(hidden_size, dtype=dtype) for _ in range(num_blocks)]) def forward(self, x: torch.Tensor) -> torch.Tensor: for block in self.blocks: @@ -56,9 +54,7 @@ def test_layerwise_offload_basic(): num_blocks = 4 # Create model - model = SimpleModelWithModuleList( - num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 - ).to(device) + model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) # Get reference output without offloading input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device) @@ -72,22 +68,18 @@ def test_layerwise_offload_basic(): for block in model.blocks: for param in block.parameters(): # Parameters should be placeholder tensors (empty) - assert param.numel() == 0, ( - "Parameters should be offloaded (empty tensors)" - ) + assert param.numel() == 0, ("Parameters should be offloaded (empty tensors)") # Run forward pass with offloading with torch.no_grad(): offloaded_output = model(input_tensor.clone()) # Check output correctness - assert torch.allclose( - reference_output, offloaded_output, rtol=1e-4, atol=1e-5 - ), "Output with offloading should match reference output" + assert torch.allclose(reference_output, offloaded_output, rtol=1e-4, + atol=1e-5), "Output with offloading should match reference output" print( - f" Layerwise offload basic test passed: max diff = {(reference_output - offloaded_output).abs().max().item()}" - ) + f" Layerwise offload basic test passed: max diff = {(reference_output - offloaded_output).abs().max().item()}") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @@ -100,14 +92,10 @@ def test_layerwise_offload_bf16(): num_blocks = 3 # Create model - model = SimpleModelWithModuleList( - num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.bfloat16 - ).to(device) + model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.bfloat16).to(device) # Get reference output without offloading - input_tensor = torch.randn( - batch_size, seq_len, hidden_size, device=device, dtype=torch.bfloat16 - ) + input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device, dtype=torch.bfloat16) with torch.no_grad(): reference_output = model(input_tensor.clone()) @@ -119,13 +107,11 @@ def test_layerwise_offload_bf16(): offloaded_output = model(input_tensor.clone()) # Check output correctness (looser tolerance for bf16) - assert torch.allclose( - reference_output, offloaded_output, rtol=1e-2, atol=1e-3 - ), "Output with offloading should match reference output for bf16" + assert torch.allclose(reference_output, offloaded_output, rtol=1e-2, + atol=1e-3), "Output with offloading should match reference output for bf16" print( - f" Layerwise offload bf16 test passed: max diff = {(reference_output - offloaded_output).abs().max().item()}" - ) + f" Layerwise offload bf16 test passed: max diff = {(reference_output - offloaded_output).abs().max().item()}") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @@ -139,17 +125,13 @@ def test_layerwise_offload_multiple_forward_passes(): num_iterations = 5 # Create model - model = SimpleModelWithModuleList( - num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 - ).to(device) + model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) # Get reference outputs without offloading torch.manual_seed(42) reference_outputs = [] for i in range(num_iterations): - input_tensor = torch.randn( - batch_size, seq_len, hidden_size, device=device - ) + input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device) with torch.no_grad(): reference_outputs.append(model(input_tensor.clone())) @@ -159,20 +141,15 @@ def test_layerwise_offload_multiple_forward_passes(): # Run multiple forward passes with offloading torch.manual_seed(42) for i in range(num_iterations): - input_tensor = torch.randn( - batch_size, seq_len, hidden_size, device=device - ) + input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device) with torch.no_grad(): offloaded_output = model(input_tensor.clone()) # Check output correctness for each iteration - assert torch.allclose( - reference_outputs[i], offloaded_output, rtol=1e-4, atol=1e-5 - ), f"Output mismatch in iteration {i}" + assert torch.allclose(reference_outputs[i], offloaded_output, rtol=1e-4, + atol=1e-5), f"Output mismatch in iteration {i}" - print( - f" Multiple forward passes test passed for {num_iterations} iterations" - ) + print(f" Multiple forward passes test passed for {num_iterations} iterations") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @@ -183,9 +160,7 @@ def test_layerwise_offload_parameter_integrity(): num_blocks = 2 # Create model - model = SimpleModelWithModuleList( - num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 - ).to(device) + model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) # Store original parameter values per block original_params_per_block = [] @@ -202,26 +177,19 @@ def test_layerwise_offload_parameter_integrity(): for block_idx, block in enumerate(model.blocks): manager = ModuleHookManager.get_from(block) assert manager is not None, "Hook manager should be attached to blocks" - hook: LayerwiseOffloadHook | None = manager.get_forward_hook( - "LayerwiseOffloadHook" - ) + hook: LayerwiseOffloadHook | None = manager.get_forward_hook("LayerwiseOffloadHook") assert hook is not None, "LayerwiseOffloadHook should be registered" # Verify parameters are stored in CPU state = hook.state - assert len(state.cpu_named_parameters) > 0, ( - "CPU parameters should be stored" - ) + assert len(state.cpu_named_parameters) > 0, ("CPU parameters should be stored") # Verify CPU parameters match original values for this block original_params = original_params_per_block[block_idx] for name, cpu_param in state.cpu_named_parameters.items(): - assert name in original_params, ( - f"CPU parameter {name} not found in original params for block {block_idx}" - ) - assert torch.allclose( - cpu_param.cpu(), original_params[name].cpu(), rtol=1e-5, atol=1e-6 - ), f"CPU parameter {name} should match original for block {block_idx}" + assert name in original_params, (f"CPU parameter {name} not found in original params for block {block_idx}") + assert torch.allclose(cpu_param.cpu(), original_params[name].cpu(), rtol=1e-5, + atol=1e-6), f"CPU parameter {name} should match original for block {block_idx}" print(" Parameter integrity test passed") @@ -231,6 +199,7 @@ def test_layerwise_offload_no_modulelist_error(): """Test that enabling offload on a model without ModuleList raises an error.""" class ModelWithoutModuleList(nn.Module): + def __init__(self): super().__init__() self.linear = nn.Linear(128, 128) @@ -241,8 +210,8 @@ def test_layerwise_offload_no_modulelist_error(): model = ModelWithoutModuleList().to("cuda") with pytest.raises( - ValueError, - match="No nn.ModuleList found in the model for layerwise offloading", + ValueError, + match="No nn.ModuleList found in the model for layerwise offloading", ): enable_layerwise_offload(model) @@ -259,9 +228,7 @@ def test_layerwise_offload_memory_reduction(): seq_len = 64 # Create model - model = SimpleModelWithModuleList( - num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 - ).to(device) + model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) # Measure initial GPU memory torch.cuda.empty_cache() @@ -274,9 +241,7 @@ def test_layerwise_offload_memory_reduction(): memory_without_offload = torch.cuda.max_memory_allocated() # Reset model - model = SimpleModelWithModuleList( - num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 - ).to(device) + model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) # Enable offloading enable_layerwise_offload(model) @@ -293,15 +258,12 @@ def test_layerwise_offload_memory_reduction(): # Memory with offload should be less (parameters are offloaded) # Note: This is a weak check as memory usage depends on many factors - print( - f"Memory without offload: {memory_without_offload / 1024**2:.2f} MB, " - f"with offload: {memory_with_offload / 1024**2:.2f} MB" - ) + print(f"Memory without offload: {memory_without_offload / 1024**2:.2f} MB, " + f"with offload: {memory_with_offload / 1024**2:.2f} MB") # We expect some reduction but this test is more informational assert memory_with_offload < memory_without_offload * 1.5, ( - "Memory usage should not increase significantly with offloading" - ) + "Memory usage should not increase significantly with offloading") print(" Memory reduction test passed") @@ -316,9 +278,7 @@ def test_layerwise_offload_gradient_disabled(): num_blocks = 3 # Create model - model = SimpleModelWithModuleList( - num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 - ).to(device) + model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) model.eval() # Enable layerwise offloading @@ -348,14 +308,11 @@ def test_layerwise_offload_different_batch_sizes(): seq_len = 16 num_blocks = 3 - model = SimpleModelWithModuleList( - num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 - ).to(device) + model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) # Get reference model without offloading - reference_model = SimpleModelWithModuleList( - num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 - ).to(device) + reference_model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, + dtype=torch.float32).to(device) reference_model.load_state_dict(model.state_dict()) # Enable layerwise offloading @@ -363,17 +320,14 @@ def test_layerwise_offload_different_batch_sizes(): # Test different batch sizes for batch_size in [1, 2, 4, 8]: - input_tensor = torch.randn( - batch_size, seq_len, hidden_size, device=device - ) + input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device) with torch.no_grad(): reference_output = reference_model(input_tensor.clone()) offloaded_output = model(input_tensor.clone()) - assert torch.allclose( - reference_output, offloaded_output, rtol=1e-4, atol=1e-5 - ), f"Output mismatch for batch_size={batch_size}" + assert torch.allclose(reference_output, offloaded_output, rtol=1e-4, + atol=1e-5), f"Output mismatch for batch_size={batch_size}" print(" Different batch sizes test passed") diff --git a/fastvideo/tests/inference/bsa/test_bsa_inference.py b/fastvideo/tests/inference/bsa/test_bsa_inference.py index e52842d8b..199b50f45 100644 --- a/fastvideo/tests/inference/bsa/test_bsa_inference.py +++ b/fastvideo/tests/inference/bsa/test_bsa_inference.py @@ -69,16 +69,12 @@ def test_inference_bsa(): }, } - with tempfile.NamedTemporaryFile( - mode="w", suffix=".json", delete=False) as f: + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: json.dump(config, f) config_path = f.name try: - cmd = [ - sys.executable, "-m", "fastvideo.entrypoints.cli.main", - "generate", "--config", config_path - ] + cmd = [sys.executable, "-m", "fastvideo.entrypoints.cli.main", "generate", "--config", config_path] subprocess.run(cmd, check=True) finally: os.unlink(config_path) diff --git a/fastvideo/tests/inference/lora/test_lora_inference_similarity.py b/fastvideo/tests/inference/lora/test_lora_inference_similarity.py index ef6d05041..c5c7753f5 100644 --- a/fastvideo/tests/inference/lora/test_lora_inference_similarity.py +++ b/fastvideo/tests/inference/lora/test_lora_inference_similarity.py @@ -16,6 +16,7 @@ from torch.testing import assert_close from torch.distributed.tensor import DTensor from fastvideo.worker import MultiprocExecutor import torch + logger = init_logger(__name__) os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29500" @@ -32,52 +33,56 @@ WAN_LORA_PARAMS = { "flow_shift": 3.0, "seed": 42, "fps": 24, - "neg_prompt": "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", - "text-encoder-precision": ("fp32",), + "neg_prompt": + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + "text-encoder-precision": ("fp32", ), "dit_cpu_offload": True, } # LoRA configurations for testing -LORA_CONFIGS = [ - { - "lora_path": "benjamin-paine/steamboat-willie-1.3b", - "lora_nickname": "steamboat", - "prompt": "steamboat willie style, golden era animation, close-up of a short fluffy monster kneeling beside a melting red candle. the mood is one of wonder and curiosity, as the monster gazes at the flame with wide eyes and open mouth. Its pose and expression convey a sense of innocence and playfulness, as if it is exploring the world around it for the first time. The use of warm colors and dramatic lighting further enhances the cozy atmosphere of the image.", - "negative_prompt": "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", - "ssim_threshold": 0.79 - }, - { - "lora_path": "motimalu/wan-flat-color-1.3b-v2", - "lora_nickname": "flat_color", - "prompt": "flat color, no lineart, blending, negative space, artist:[john kafka|ponsuke kaikai|hara id 21|yoneyama mai|fuzichoco], 1girl, sakura miko, pink hair, cowboy shot, white shirt, floral print, off shoulder, outdoors, cherry blossom, tree shade, wariza, looking up, falling petals, half-closed eyes, white sky, clouds, live2d animation, upper body, high quality cinematic video of a woman sitting under a sakura tree. Dreamy and lonely, the camera close-ups on the face of the woman as she turns towards the viewer. The Camera is steady, This is a cowboy shot. The animation is smooth and fluid.", - "negative_prompt": "bad quality video,色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", - "ssim_threshold": 0.79 - } - # TODO: Add a LoRA with lora_alpha values to test alpha scaling - # - # Context: This change is mainly for an in-progress ticket porting over LongCat-Video, - # where they used an alpha value that is two times smaller than their rank. This fix - # ensures that LoRA weights are correctly scaled by the alpha/rank ratio when merged. - # - # Issue: Currently, we cannot add a test for LoRA adapters with alpha values because: - # - The existing public LoRAs for Wan-AI/Wan2.1-T2V-1.3B-Diffusers don't store lora_alpha - # - No publicly available LoRA for this model includes lora_alpha tensors in their weights - # - This is why the alpha/rank scaling bug wasn't caught by existing tests - # - # The fix has been validated with: - # - LongCat-Video distilled LoRA (which includes alpha values) - # - Manual testing shows correct alpha/rank scaling behavior - # - Backward compatibility confirmed with LoRAs without alpha values - # - # Future work: - # - Add a synthetic LoRA test fixture with alpha values when feasible - # - Or wait for public Wan LoRAs with alpha to become available -] +LORA_CONFIGS = [{ + "lora_path": "benjamin-paine/steamboat-willie-1.3b", + "lora_nickname": "steamboat", + "prompt": + "steamboat willie style, golden era animation, close-up of a short fluffy monster kneeling beside a melting red candle. the mood is one of wonder and curiosity, as the monster gazes at the flame with wide eyes and open mouth. Its pose and expression convey a sense of innocence and playfulness, as if it is exploring the world around it for the first time. The use of warm colors and dramatic lighting further enhances the cozy atmosphere of the image.", + "negative_prompt": + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + "ssim_threshold": 0.79 +}, { + "lora_path": "motimalu/wan-flat-color-1.3b-v2", + "lora_nickname": "flat_color", + "prompt": + "flat color, no lineart, blending, negative space, artist:[john kafka|ponsuke kaikai|hara id 21|yoneyama mai|fuzichoco], 1girl, sakura miko, pink hair, cowboy shot, white shirt, floral print, off shoulder, outdoors, cherry blossom, tree shade, wariza, looking up, falling petals, half-closed eyes, white sky, clouds, live2d animation, upper body, high quality cinematic video of a woman sitting under a sakura tree. Dreamy and lonely, the camera close-ups on the face of the woman as she turns towards the viewer. The Camera is steady, This is a cowboy shot. The animation is smooth and fluid.", + "negative_prompt": + "bad quality video,色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + "ssim_threshold": 0.79 +} + # TODO: Add a LoRA with lora_alpha values to test alpha scaling + # + # Context: This change is mainly for an in-progress ticket porting over LongCat-Video, + # where they used an alpha value that is two times smaller than their rank. This fix + # ensures that LoRA weights are correctly scaled by the alpha/rank ratio when merged. + # + # Issue: Currently, we cannot add a test for LoRA adapters with alpha values because: + # - The existing public LoRAs for Wan-AI/Wan2.1-T2V-1.3B-Diffusers don't store lora_alpha + # - No publicly available LoRA for this model includes lora_alpha tensors in their weights + # - This is why the alpha/rank scaling bug wasn't caught by existing tests + # + # The fix has been validated with: + # - LongCat-Video distilled LoRA (which includes alpha values) + # - Manual testing shows correct alpha/rank scaling behavior + # - Backward compatibility confirmed with LoRAs without alpha values + # + # Future work: + # - Add a synthetic LoRA test fixture with alpha values when feasible + # - Or wait for public Wan LoRAs with alpha to become available + ] MODEL_TO_PARAMS = { "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WAN_LORA_PARAMS, } + def _sanitize_filename_component(name: str) -> str: """Sanitize filename to remove invalid characters (same logic as VideoGenerator)""" sanitized = re.sub(r'[\\/:*?"<>|]', '', name) @@ -85,9 +90,10 @@ def _sanitize_filename_component(name: str) -> str: sanitized = re.sub(r'\s+', ' ', sanitized) return sanitized or "video" + @pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys())) def test_merge_lora_weights(model_id): - lora_config = LORA_CONFIGS[0] # test only one + lora_config = LORA_CONFIGS[0] # test only one hf_pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) hf_pipe.enable_model_cpu_offload() @@ -119,9 +125,11 @@ def test_merge_lora_weights(model_id): if "base_layer" not in key: continue hf_param = hf_state_dict[key] - custom_param = custom_state_dict[key].to_local() if isinstance(custom_state_dict[key], DTensor) else custom_state_dict[key] + custom_param = custom_state_dict[key].to_local() if isinstance(custom_state_dict[key], + DTensor) else custom_state_dict[key] assert_close(hf_param, custom_param, atol=7e-4, rtol=7e-4) + @pytest.mark.parametrize("ATTENTION_BACKEND", ["TORCH_SDPA"]) @pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys())) def test_lora_inference_similarity(ATTENTION_BACKEND, model_id): @@ -175,30 +183,28 @@ def test_lora_inference_similarity(ATTENTION_BACKEND, model_id): generator.generate_video(prompt, **generation_kwargs) - assert os.path.exists( - generated_video_path), f"Output video was not generated at {generated_video_path}" + assert os.path.exists(generated_video_path), f"Output video was not generated at {generated_video_path}" reference_folder = os.path.join(script_dir, 'L40S_reference_videos', model_id.split('/')[-1], ATTENTION_BACKEND) - + if not os.path.exists(reference_folder): logger.error("Reference folder missing") - raise FileNotFoundError( - f"Reference video folder does not exist: {reference_folder}") + raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") # Find the matching reference video - try exact match first, then fuzzy match # The reference might have different sanitization (e.g., trailing spaces) reference_video_name = None unsanitized_prefix = f"{lora_path.split('/')[-1]}_{prompt[:50]}" - + for filename in os.listdir(reference_folder): if not filename.endswith('.mp4'): continue - + # Try exact match with sanitized name if filename.startswith(output_video_name): reference_video_name = filename break - + # Try match with unsanitized prefix (for legacy reference videos) # Remove .mp4 and compare the base names after sanitization base_filename = filename[:-4] # Remove .mp4 @@ -207,32 +213,25 @@ def test_lora_inference_similarity(ATTENTION_BACKEND, model_id): break if not reference_video_name: - logger.error(f"Reference video not found for adapter: {lora_path} with prompt: {prompt[:50]} and backend: {ATTENTION_BACKEND}") + logger.error( + f"Reference video not found for adapter: {lora_path} with prompt: {prompt[:50]} and backend: {ATTENTION_BACKEND}" + ) raise FileNotFoundError(f"Reference video missing for adapter {lora_path}") reference_video_path = os.path.join(reference_folder, reference_video_name) - logger.info( - f"Computing SSIM between {reference_video_path} and {generated_video_path}" - ) - ssim_values = compute_video_ssim_torchvision(reference_video_path, - generated_video_path, - use_ms_ssim=True) + logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") + ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) mean_ssim = ssim_values[0] logger.info(f"SSIM mean value: {mean_ssim}") logger.info(f"Writing SSIM results to directory: {output_dir}") - success = write_ssim_results(output_dir, ssim_values, reference_video_path, - generated_video_path, num_inference_steps, - prompt) + success = write_ssim_results(output_dir, ssim_values, reference_video_path, generated_video_path, + num_inference_steps, prompt) if not success: logger.error("Failed to write SSIM results to file") min_acceptable_ssim = lora_config["ssim_threshold"] - assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for adapter {lora_config['lora_path']}" - - - - + assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for adapter {lora_config['lora_path']}" diff --git a/fastvideo/tests/inference/test_cfg_gating.py b/fastvideo/tests/inference/test_cfg_gating.py index 07d833db6..1f9f5e495 100644 --- a/fastvideo/tests/inference/test_cfg_gating.py +++ b/fastvideo/tests/inference/test_cfg_gating.py @@ -8,6 +8,7 @@ import torch class RecordingLogger: + def __init__(self): self.infos = [] self.warnings = [] @@ -20,6 +21,7 @@ class RecordingLogger: class NullProgressBar: + def __enter__(self): return self @@ -56,8 +58,9 @@ class TinyDenoiser(torch.nn.Module): self.calls.append("uncond" if is_uncond else "cond") prompt_term = prompt_value.reshape((1, ) * latent_model_input.ndim) - timestep_term = timestep.to(latent_model_input.device, torch.float32).reshape( - latent_model_input.shape[0], *([1] * (latent_model_input.ndim - 1))) + timestep_term = timestep.to(latent_model_input.device, + torch.float32).reshape(latent_model_input.shape[0], + *([1] * (latent_model_input.ndim - 1))) return latent_model_input.float() * 0.2 + prompt_term * 0.5 + timestep_term * 0.001 diff --git a/fastvideo/tests/inference/vmoba/test_vmoba_inference.py b/fastvideo/tests/inference/vmoba/test_vmoba_inference.py index 694b76a4b..333a3a9f8 100644 --- a/fastvideo/tests/inference/vmoba/test_vmoba_inference.py +++ b/fastvideo/tests/inference/vmoba/test_vmoba_inference.py @@ -70,8 +70,7 @@ def test_inference_vmoba(): }, } - with tempfile.NamedTemporaryFile( - mode="w", suffix=".json", delete=False) as f: + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: json.dump(config, f) config_path = f.name diff --git a/fastvideo/tests/layers/test_rmsnorm_forward_dispatch.py b/fastvideo/tests/layers/test_rmsnorm_forward_dispatch.py index de9a2d113..8b22c5bf8 100644 --- a/fastvideo/tests/layers/test_rmsnorm_forward_dispatch.py +++ b/fastvideo/tests/layers/test_rmsnorm_forward_dispatch.py @@ -50,17 +50,13 @@ def _run_torchrun(script_path: Path, mode: str, output_path: Path) -> None: timeout=120, ) except subprocess.TimeoutExpired as error: - raise RuntimeError( - f"{mode} worker timed out after 120 seconds\n" - f"STDOUT:\n{error.stdout}\n" - f"STDERR:\n{error.stderr}" - ) from error + raise RuntimeError(f"{mode} worker timed out after 120 seconds\n" + f"STDOUT:\n{error.stdout}\n" + f"STDERR:\n{error.stderr}") from error if process.returncode != 0: - raise RuntimeError( - f"{mode} worker failed with code {process.returncode}\n" - f"STDOUT:\n{process.stdout}\n" - f"STDERR:\n{process.stderr}" - ) + raise RuntimeError(f"{mode} worker failed with code {process.returncode}\n" + f"STDOUT:\n{process.stdout}\n" + f"STDERR:\n{process.stderr}") def _summarize_tensor(tensor: torch.Tensor | Any) -> dict[str, Any]: @@ -75,10 +71,10 @@ def _summarize_tensor(tensor: torch.Tensor | Any) -> dict[str, Any]: def _run_worker(mode: str, output_path: Path) -> None: if mode not in { - "module_no_offload", - "direct_no_offload", - "module_cpu_offload", - "direct_cpu_offload", + "module_no_offload", + "direct_no_offload", + "module_cpu_offload", + "direct_cpu_offload", }: raise ValueError(f"Unsupported mode: {mode}") @@ -91,7 +87,7 @@ def _run_worker(mode: str, output_path: Path) -> None: torch.manual_seed(SEED + rank) try: - mesh = init_device_mesh("cuda", (world_size,)) + mesh = init_device_mesh("cuda", (world_size, )) norm = RMSNorm(HIDDEN_SIZE, eps=1e-6, has_weight=True).to(device) with torch.no_grad(): norm.weight.fill_(1.0) @@ -179,8 +175,7 @@ def test_no_direct_forward_native_calls_in_models() -> None: hooks (issue #1379); model code must use module dispatch instead.""" models_dir = REPO_ROOT / "fastvideo" / "models" offenders = [ - str(path.relative_to(REPO_ROOT)) - for path in sorted(models_dir.rglob("*.py")) + str(path.relative_to(REPO_ROOT)) for path in sorted(models_dir.rglob("*.py")) if ".forward_native(" in path.read_text(encoding="utf-8") ] assert not offenders, f"Replace .forward_native(...) with module dispatch in: {offenders}" diff --git a/fastvideo/tests/lora_extraction/test_lora_extraction.py b/fastvideo/tests/lora_extraction/test_lora_extraction.py index 02bb1682c..51e35b452 100644 --- a/fastvideo/tests/lora_extraction/test_lora_extraction.py +++ b/fastvideo/tests/lora_extraction/test_lora_extraction.py @@ -16,13 +16,13 @@ from verify_lora import main as verify_lora_main def test_lora_extraction_pipeline(): """Test end-to-end LoRA extraction workflow.""" import tempfile - + # Use temp directory for outputs to avoid polluting repo with tempfile.TemporaryDirectory() as tmpdir: tmpdir_path = Path(tmpdir) adapter_path = tmpdir_path / "adapter_r16.safetensors" merged_dir = tmpdir_path / "merged_r16" - + # 1. Extract rank-16 adapter print("\nExtracting rank-16 adapter") extract_lora_adapter( @@ -32,7 +32,7 @@ def test_lora_extraction_pipeline(): rank=16, ) assert adapter_path.exists(), "Adapter file was not created" - + # 2. Merge adapter print("\nMerging adapter") merge_lora( @@ -42,7 +42,7 @@ def test_lora_extraction_pipeline(): output=str(merged_dir), ) assert merged_dir.exists(), "Merged model directory was not created" - + # 3. Verify numerical accuracy print("\nVerifying merged model") # verify_lora uses sys.argv, so we need to mock it @@ -50,11 +50,13 @@ def test_lora_extraction_pipeline(): try: sys.argv = [ "verify_lora.py", - "--merged", str(merged_dir), - "--ft", "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", + "--merged", + str(merged_dir), + "--ft", + "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", ] verify_lora_main() finally: sys.argv = old_argv - + print("\nLoRA extraction pipeline test PASSED") diff --git a/fastvideo/tests/modal/kernel_build_cache.py b/fastvideo/tests/modal/kernel_build_cache.py index 936a7b5ad..26eb19cf5 100644 --- a/fastvideo/tests/modal/kernel_build_cache.py +++ b/fastvideo/tests/modal/kernel_build_cache.py @@ -22,7 +22,6 @@ import zipfile from email.parser import Parser from pathlib import Path - CACHE_SCHEMA_VERSION = 3 DEFAULT_PREBUILT_INFO_PATH = "/opt/fastvideo-kernel-prebuilt" KERNEL_RELATIVE_DIR = "fastvideo-kernel" @@ -39,11 +38,7 @@ def _log(message: str) -> None: print(f"[fastvideo-kernel-cache] {message}", flush=True) -def _run(args: list[str], - *, - cwd: Path | None = None, - env: dict[str, str] | None = None, - capture: bool = False) -> str: +def _run(args: list[str], *, cwd: Path | None = None, env: dict[str, str] | None = None, capture: bool = False) -> str: _log("$ " + " ".join(str(arg) for arg in args)) result = subprocess.run( args, @@ -82,8 +77,8 @@ def _hash_directory(root: Path) -> str: hasher = hashlib.sha256() skip_dirs = {".git", ".mypy_cache", ".pytest_cache", "__pycache__", "build", "dist"} for current_root, dirnames, filenames in os.walk(root): - dirnames[:] = sorted( - dirname for dirname in dirnames if dirname not in skip_dirs and not dirname.endswith(".egg-info")) + dirnames[:] = sorted(dirname for dirname in dirnames + if dirname not in skip_dirs and not dirname.endswith(".egg-info")) for filename in sorted(filenames): path = Path(current_root) / filename hasher.update(path.relative_to(root).as_posix().encode("utf-8")) @@ -100,8 +95,7 @@ def _kernel_source_hash(repo_root: Path) -> dict[str, str]: try: tree_hash = _run(["git", "rev-parse", f"HEAD:{KERNEL_RELATIVE_DIR}"], cwd=repo_root, capture=True) diff = _run_optional(["git", "diff", "--binary", "--", KERNEL_RELATIVE_DIR], cwd=repo_root) - cached_diff = _run_optional(["git", "diff", "--cached", "--binary", "--", KERNEL_RELATIVE_DIR], - cwd=repo_root) + cached_diff = _run_optional(["git", "diff", "--cached", "--binary", "--", KERNEL_RELATIVE_DIR], cwd=repo_root) submodules = _run_optional( [ "git", @@ -268,11 +262,8 @@ def _build_metadata(repo_root: Path) -> dict[str, object]: def _find_wheel(directory: Path) -> Path: - candidates = sorted( - path - for pattern in ("fastvideo_kernel-*.whl", "fastvideo-kernel-*.whl") - for path in directory.glob(pattern) - ) + candidates = sorted(path for pattern in ("fastvideo_kernel-*.whl", "fastvideo-kernel-*.whl") + for path in directory.glob(pattern)) if not candidates: raise RuntimeError(f"No fastvideo-kernel wheel found in {directory}") return candidates[-1] @@ -427,8 +418,7 @@ def _store_cache_entry(cache_root: Path, metadata: dict[str, object], wheel: Pat "created_at_utc": _utc_now_isoformat(), "artifact": _wheel_artifact(stored_wheel), } - (temp_entry / METADATA_FILE).write_text(json.dumps(stored_metadata, indent=2, sort_keys=True), - encoding="utf-8") + (temp_entry / METADATA_FILE).write_text(json.dumps(stored_metadata, indent=2, sort_keys=True), encoding="utf-8") temp_entry.rename(cache_entry) except OSError as error: shutil.rmtree(temp_entry, ignore_errors=True) diff --git a/fastvideo/tests/modal/kernel_cache_smoke.py b/fastvideo/tests/modal/kernel_cache_smoke.py index c140caab3..8305ec9fc 100644 --- a/fastvideo/tests/modal/kernel_cache_smoke.py +++ b/fastvideo/tests/modal/kernel_cache_smoke.py @@ -50,24 +50,22 @@ print(f"Using kernel cache volume: {KERNEL_CACHE_VOLUME_NAME}") kernel_cache_vol = modal.Volume.from_name(KERNEL_CACHE_VOLUME_NAME, create_if_missing=True) uv_torch_backend_override = resolve_uv_torch_backend(IMAGE_TAG) -image = ( - modal.Image.from_registry(IMAGE_REF, add_python="3.12") - .apt_install( - "cmake", - "pkg-config", - "build-essential", - "curl", - "git", - "libssl-dev", - "ffmpeg", - ) - .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable") - .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc") - .env({ - "PATH": "/root/.cargo/bin:$PATH", - **({"UV_TORCH_BACKEND": uv_torch_backend_override} if uv_torch_backend_override else {}), - }) -) +image = (modal.Image.from_registry(IMAGE_REF, add_python="3.12").apt_install( + "cmake", + "pkg-config", + "build-essential", + "curl", + "git", + "libssl-dev", + "ffmpeg", +).run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable"). + run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc").env({ + "PATH": + "/root/.cargo/bin:$PATH", + **({ + "UV_TORCH_BACKEND": uv_torch_backend_override + } if uv_torch_backend_override else {}), + })) PRODUCER_KWARGS = dict( gpu="L40S:1", diff --git a/fastvideo/tests/modal/launch_l40s_job.py b/fastvideo/tests/modal/launch_l40s_job.py index 94da57554..c86ffb88e 100644 --- a/fastvideo/tests/modal/launch_l40s_job.py +++ b/fastvideo/tests/modal/launch_l40s_job.py @@ -44,6 +44,7 @@ except ModuleNotFoundError: def resolve_uv_torch_backend(image_tag: str) -> str | None: return os.environ.get("UV_TORCH_BACKEND") + app = modal.App("fastvideo-gpu-job") REPO_DIR = "/FastVideo" @@ -67,11 +68,7 @@ print(f"Using image: {IMAGE_REF}") print(f"Using Modal volume: {MODEL_VOLUME_NAME}") model_vol = modal.Volume.from_name(MODEL_VOLUME_NAME, create_if_missing=True) -local_secrets = modal.Secret.from_dict({ - key: os.environ[key] - for key in SECRET_ENV_KEYS - if os.environ.get(key) -}) +local_secrets = modal.Secret.from_dict({key: os.environ[key] for key in SECRET_ENV_KEYS if os.environ.get(key)}) # Mutable tags inherit the registry image's baked backend, including custom # FASTVIDEO_MODAL_IMAGE overrides. Explicit CUDA tags also work with older @@ -79,8 +76,7 @@ local_secrets = modal.Secret.from_dict({ uv_torch_backend_override = resolve_uv_torch_backend(IMAGE_TAG) image = ( - modal.Image.from_registry(IMAGE_REF, add_python="3.12") - .apt_install( + modal.Image.from_registry(IMAGE_REF, add_python="3.12").apt_install( "cmake", "pkg-config", "build-essential", @@ -88,26 +84,32 @@ image = ( "git", "libssl-dev", "ffmpeg", - ) - .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable") - .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc") - .env({ - "PATH": "/root/.cargo/bin:$PATH", - "HF_HOME": "/root/data/.cache", - "TOKENIZERS_PARALLELISM": "false", - "IMAGE_VERSION": IMAGE_VERSION, - "FASTVIDEO_CONTAINER_IMAGE_REF": IMAGE_REF, - **({"UV_TORCH_BACKEND": uv_torch_backend_override} if uv_torch_backend_override else {}), - "FASTVIDEO_ATTENTION_BACKEND": os.environ.get("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN"), + ).run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable") + .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc").env({ + "PATH": + "/root/.cargo/bin:$PATH", + "HF_HOME": + "/root/data/.cache", + "TOKENIZERS_PARALLELISM": + "false", + "IMAGE_VERSION": + IMAGE_VERSION, + "FASTVIDEO_CONTAINER_IMAGE_REF": + IMAGE_REF, + **({ + "UV_TORCH_BACKEND": uv_torch_backend_override + } if uv_torch_backend_override else {}), + "FASTVIDEO_ATTENTION_BACKEND": + os.environ.get("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN"), **({ "FASTVIDEO_PERFORMANCE_PROFILE_VERSION": os.environ["FASTVIDEO_PERFORMANCE_PROFILE_VERSION"] } if os.environ.get("FASTVIDEO_PERFORMANCE_PROFILE_VERSION") else {}), # FA4 is opt-in (FASTVIDEO_FA4). Generic ad hoc jobs should follow the # product default unless a caller opts in through the local env or # --env-vars. - "FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "0"), - }) -) + "FASTVIDEO_FA4": + os.environ.get("FASTVIDEO_FA4", "0"), + })) COMMON_FUNCTION_KWARGS = dict( image=image, @@ -502,7 +504,8 @@ def _select_runner(gpu_type: str, num_gpus: int) -> Callable[..., Any]: return runners[(normalized_gpu_type, num_gpus)] except KeyError as error: supported = ", ".join(f"{gpu}:{count}" for gpu, count in sorted(runners)) - raise RuntimeError(f"Unsupported GPU request {gpu_type}:{num_gpus}. Supported requests: {supported}.") from error + raise RuntimeError( + f"Unsupported GPU request {gpu_type}:{num_gpus}. Supported requests: {supported}.") from error @app.local_entrypoint() diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index 4124271f4..fc007dd32 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -21,6 +21,7 @@ except ModuleNotFoundError: def resolve_uv_torch_backend(image_tag: str) -> str | None: return os.environ.get("UV_TORCH_BACKEND") + app = modal.App() model_vol = modal.Volume.from_name("hf-model-weights") @@ -43,20 +44,16 @@ uv_torch_backend_override = resolve_uv_torch_backend(image_tag) # budget. `image_ref` is the only env-derived input allowed here, because it # *selects* the base digest. Per-job values arrive at runtime via # `ci_env_secret` below. -image = (modal.Image.from_registry( - image_ref, add_python="3.12" -).run_commands("rm -rf /FastVideo").apt_install( - "cmake", "pkg-config", "build-essential", "curl", "libssl-dev", "ffmpeg" -).run_commands( - "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable" -).run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc").env({ - "PATH": "/root/.cargo/bin:$PATH", - "HF_REPO_ID": "FastVideo/performance-tracking", -})) +image = (modal.Image.from_registry(image_ref, add_python="3.12").run_commands("rm -rf /FastVideo").apt_install( + "cmake", "pkg-config", "build-essential", "curl", "libssl-dev", "ffmpeg").run_commands( + "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable"). + run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc").env({ + "PATH": "/root/.cargo/bin:$PATH", + "HF_REPO_ID": "FastVideo/performance-tracking", + })) -dreamverse_image = (image.run_commands( - "curl -fsSL https://deb.nodesource.com/setup_22.x | bash -" -).apt_install("nodejs").run_commands("node --version && npm --version")) +dreamverse_image = (image.run_commands("curl -fsSL https://deb.nodesource.com/setup_22.x | bash -").apt_install( + "nodejs").run_commands("node --version && npm --version")) # Per-job/per-invocation values are injected into the container environment at # RUNTIME via this secret (attached to every function below), so the image @@ -65,24 +62,34 @@ dreamverse_image = (image.run_commands( # resolver, `uv pip install`) all read os.environ at runtime, so nothing else # changes. ci_env_secret = modal.Secret.from_dict({ - "BUILDKITE_REPO": os.environ.get("BUILDKITE_REPO", ""), - "BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""), - "BUILDKITE_PULL_REQUEST": os.environ.get("BUILDKITE_PULL_REQUEST", ""), - "BUILDKITE_BRANCH": os.environ.get("BUILDKITE_BRANCH", ""), - "BUILDKITE_SOURCE": os.environ.get("BUILDKITE_SOURCE", ""), - "BUILDKITE_BUILD_URL": os.environ.get("BUILDKITE_BUILD_URL", ""), - "BUILDKITE_BUILD_ID": os.environ.get("BUILDKITE_BUILD_ID", ""), - "BUILDKITE_JOB_ID": os.environ.get("BUILDKITE_JOB_ID", ""), - "TEST_SCOPE": os.environ.get("TEST_SCOPE", ""), - "IMAGE_VERSION": image_version, - "FASTVIDEO_CONTAINER_IMAGE_REF": image_ref, + "BUILDKITE_REPO": + os.environ.get("BUILDKITE_REPO", ""), + "BUILDKITE_COMMIT": + os.environ.get("BUILDKITE_COMMIT", ""), + "BUILDKITE_PULL_REQUEST": + os.environ.get("BUILDKITE_PULL_REQUEST", ""), + "BUILDKITE_BRANCH": + os.environ.get("BUILDKITE_BRANCH", ""), + "BUILDKITE_SOURCE": + os.environ.get("BUILDKITE_SOURCE", ""), + "BUILDKITE_BUILD_URL": + os.environ.get("BUILDKITE_BUILD_URL", ""), + "BUILDKITE_BUILD_ID": + os.environ.get("BUILDKITE_BUILD_ID", ""), + "BUILDKITE_JOB_ID": + os.environ.get("BUILDKITE_JOB_ID", ""), + "TEST_SCOPE": + os.environ.get("TEST_SCOPE", ""), + "IMAGE_VERSION": + image_version, + "FASTVIDEO_CONTAINER_IMAGE_REF": + image_ref, **{ key: os.environ[key] for key in ( "FASTVIDEO_ATTENTION_BACKEND", "FASTVIDEO_PERFORMANCE_PROFILE_VERSION", - ) - if os.environ.get(key) + ) if os.environ.get(key) }, **({ "UV_TORCH_BACKEND": uv_torch_backend_override @@ -91,19 +98,15 @@ ci_env_secret = modal.Secret.from_dict({ # inference/perf parity; model-load and training lanes that do not exercise # FA4 explicitly set FASTVIDEO_FA4=0 in their command strings below. # Caller override wins. - "FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1"), + "FASTVIDEO_FA4": + os.environ.get("FASTVIDEO_FA4", "1"), }) -hf_secret = modal.Secret.from_dict( - {"HF_API_KEY": os.environ.get("HF_API_KEY", "")}) -wandb_secret = modal.Secret.from_dict( - {"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")}) +hf_secret = modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")}) +wandb_secret = modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")}) -def _run_git_with_retries(command: list[str], - *, - cwd: str, - cleanup_path: str | None = None) -> None: +def _run_git_with_retries(command: list[str], *, cwd: str, cleanup_path: str | None = None) -> None: last_returncode = 1 for attempt in range(1, 4): if cleanup_path is not None: @@ -116,21 +119,15 @@ def _run_git_with_retries(command: list[str], last_returncode = result.returncode if attempt < 3: sleep_seconds = 5 * attempt - print( - f"Git command failed (attempt {attempt}/3, exit {last_returncode}); " - f"retrying in {sleep_seconds}s", - flush=True) + print(f"Git command failed (attempt {attempt}/3, exit {last_returncode}); " + f"retrying in {sleep_seconds}s", + flush=True) time.sleep(sleep_seconds) - raise RuntimeError( - f"Git command failed after 3 attempts with exit code {last_returncode}: " - + " ".join(command)) + raise RuntimeError(f"Git command failed after 3 attempts with exit code {last_returncode}: " + " ".join(command)) -def _checkout_repository(git_repo: str, - git_commit: str, - pr_number: str | None, - repo_root: str = "/FastVideo") -> None: +def _checkout_repository(git_repo: str, git_commit: str, pr_number: str | None, repo_root: str = "/FastVideo") -> None: if not git_repo or git_repo.startswith("-"): raise RuntimeError("BUILDKITE_REPO must be a non-empty repository URL.") @@ -138,18 +135,14 @@ def _checkout_repository(git_repo: str, try: pr_id = int(pr_number) except ValueError as error: - raise RuntimeError( - f"Invalid BUILDKITE_PULL_REQUEST value: {pr_number}") from error + raise RuntimeError(f"Invalid BUILDKITE_PULL_REQUEST value: {pr_number}") from error if pr_id <= 0: - raise RuntimeError( - f"Invalid BUILDKITE_PULL_REQUEST value: {pr_number}") + raise RuntimeError(f"Invalid BUILDKITE_PULL_REQUEST value: {pr_number}") target = f"refs/pull/{pr_id}/head" print(f"Using PR ref for checkout: {target}") else: - if not git_commit or re.fullmatch(r"[0-9a-fA-F]{7,64}", - git_commit) is None: - raise RuntimeError( - f"Invalid BUILDKITE_COMMIT value: {git_commit}") + if not git_commit or re.fullmatch(r"[0-9a-fA-F]{7,64}", git_commit) is None: + raise RuntimeError(f"Invalid BUILDKITE_COMMIT value: {git_commit}") target = git_commit print(f"Using direct commit checkout: {target}") @@ -166,27 +159,21 @@ def _checkout_repository(git_repo: str, git_repo, repo_root, ] - _run_git_with_retries(clone_command, - cwd="/", - cleanup_path=repo_root) + _run_git_with_retries(clone_command, cwd="/", cleanup_path=repo_root) git_prefix = ["git", "-c", "http.version=HTTP/1.1"] - _run_git_with_retries( - git_prefix + [ - "fetch", - "--prune", - "--no-tags", - "--depth=1", - "--filter=blob:none", - "origin", - target, - ], - cwd=repo_root) - _run_git_with_retries( - git_prefix + ["checkout", "--detach", "FETCH_HEAD"], cwd=repo_root) - _run_git_with_retries( - git_prefix + ["submodule", "update", "--init", "--recursive"], - cwd=repo_root) + _run_git_with_retries(git_prefix + [ + "fetch", + "--prune", + "--no-tags", + "--depth=1", + "--filter=blob:none", + "origin", + target, + ], + cwd=repo_root) + _run_git_with_retries(git_prefix + ["checkout", "--detach", "FETCH_HEAD"], cwd=repo_root) + _run_git_with_retries(git_prefix + ["submodule", "update", "--init", "--recursive"], cwd=repo_root) def run_test(pytest_command: str): @@ -194,9 +181,7 @@ def run_test(pytest_command: str): run_test_command(pytest_command, build_kernel=True) -def run_test_command(test_command: str, - build_kernel: bool, - install_command: str = 'uv pip install -e ".[test]"'): +def run_test_command(test_command: str, build_kernel: bool, install_command: str = 'uv pip install -e ".[test]"'): """Helper function to run a test suite with custom test command. Most FastVideo CI suites need the custom kernel build. App-level tests like @@ -234,14 +219,9 @@ def run_test_command(test_command: str, setup_steps.append("python fastvideo/tests/modal/kernel_build_cache.py install") setup_command = " &&\n ".join(setup_steps) - setup_result = subprocess.run(["/bin/bash", "-c", setup_command], - stdout=sys.stdout, - stderr=sys.stderr, - check=False) + setup_result = subprocess.run(["/bin/bash", "-c", setup_command], stdout=sys.stdout, stderr=sys.stderr, check=False) if setup_result.returncode != 0: - raise RuntimeError( - f"Setup command failed with exit code {setup_result.returncode}") - + raise RuntimeError(f"Setup command failed with exit code {setup_result.returncode}") command = " &&\n ".join([ "source $HOME/.local/bin/env", @@ -249,15 +229,12 @@ def run_test_command(test_command: str, "cd /FastVideo", test_command, ]) - result = subprocess.run(["/bin/bash", "-c", command], - stdout=sys.stdout, - stderr=sys.stderr, - check=False) + result = subprocess.run(["/bin/bash", "-c", command], stdout=sys.stdout, stderr=sys.stderr, check=False) # Modal containers crash on sys.exit(0); raise on failure, return on success. if result.returncode != 0: - raise RuntimeError( - f"Test command failed with exit code {result.returncode}") + raise RuntimeError(f"Test command failed with exit code {result.returncode}") + @app.function(gpu="H100:1", image=image, @@ -277,8 +254,7 @@ def run_encoder_tests(): volumes={"/root/data": model_vol}) def run_vae_tests(): run_test( - "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/vaes -vs" - ) + "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/vaes -vs") @app.function(gpu="L40S:1", @@ -302,10 +278,8 @@ def run_golden_gate_tests(): secrets=[hf_secret, ci_env_secret], volumes={"/root/data": model_vol}) def run_transformer_tests(): - run_test( - "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && " - "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/transformers -vs" - ) + run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && " + "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/transformers -vs") @app.function(gpu="L40S:4", @@ -316,10 +290,8 @@ def run_transformer_tests(): secrets=[wandb_secret, ci_env_secret], volumes={"/root/data": model_vol}) def run_training_tests(): - run_test( - "export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && " - "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/Vanilla -srP" - ) + run_test("export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && " + "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/Vanilla -srP") @app.function(gpu="L40S:2", @@ -330,20 +302,13 @@ def run_training_tests(): secrets=[wandb_secret, ci_env_secret], volumes={"/root/data": model_vol}) def run_training_lora_tests(): - run_test( - "export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && " - "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP" - ) + run_test("export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && " + "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP") -@app.function(gpu="H100!:2", - image=image, - timeout=900, - secrets=[wandb_secret, ci_env_secret]) +@app.function(gpu="H100!:2", image=image, timeout=900, secrets=[wandb_secret, ci_env_secret]) def run_training_tests_VSA(): - run_test( - "wandb login $WANDB_API_KEY && FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/VSA -srP" - ) + run_test("wandb login $WANDB_API_KEY && FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/VSA -srP") @app.function(gpu="H100:1", image=image, timeout=900, secrets=[ci_env_secret]) @@ -368,53 +333,39 @@ def run_inference_tests_vmoba(): @app.function(gpu="L40S:1", image=image, timeout=1200, secrets=[ci_env_secret]) def run_inference_lora_tests(): - run_test( - "pytest ./fastvideo/tests/inference/lora/test_lora_inference_similarity.py -vs" - ) + run_test("pytest ./fastvideo/tests/inference/lora/test_lora_inference_similarity.py -vs") @app.function(gpu="L40S:2", image=image, timeout=900, secrets=[ci_env_secret]) def run_distill_dmd_tests(): - run_test( - "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs") + run_test("FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs") -@app.function(gpu="L40S:2", - image=image, - timeout=900, - secrets=[wandb_secret, ci_env_secret]) +@app.function(gpu="L40S:2", image=image, timeout=900, secrets=[wandb_secret, ci_env_secret]) def run_self_forcing_tests(): - run_test( - "wandb login $WANDB_API_KEY && " - "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs" - ) + run_test("wandb login $WANDB_API_KEY && " + "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs") @app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret]) def run_unit_test(): - run_test( - "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ " - "./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ " - "./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/worker/ " - "./fastvideo/tests/training/test_trackers.py " - "./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py " - "./fastvideo/tests/modal/test_kernel_build_cache.py ./fastvideo/tests/modal/test_pr_test.py " - "./fastvideo/tests/modal/test_ssim_test.py " - "--ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py " - "--ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs" - ) + run_test("pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ " + "./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ " + "./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/worker/ " + "./fastvideo/tests/training/test_trackers.py " + "./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py " + "./fastvideo/tests/modal/test_kernel_build_cache.py ./fastvideo/tests/modal/test_pr_test.py " + "./fastvideo/tests/modal/test_ssim_test.py " + "--ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py " + "--ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs") # TODO: David: GPU only used to resolve import time requirement (not needed for this test). Maybe make those imports lazy? -@app.function(gpu="L40S:1", - image=dreamverse_image, - timeout=1800, - secrets=[ci_env_secret]) +@app.function(gpu="L40S:1", image=dreamverse_image, timeout=1800, secrets=[ci_env_secret]) def run_dreamverse_app_tests(): - run_test_command( - install_command="", - build_kernel=False, - test_command=""" + run_test_command(install_command="", + build_kernel=False, + test_command=""" uv pip install -e ".[test,dreamverse]" && export PYTHONPATH=/FastVideo/apps/dreamverse:$PYTHONPATH && pytest apps/dreamverse/dreamverse/tests -q && @@ -453,10 +404,8 @@ def run_dreamverse_app_tests(): secrets=[hf_secret, ci_env_secret], volumes={"/root/data": model_vol}) def run_train_framework_tests(): - run_test( - "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && " - "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs" - ) + run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && " + "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs") @app.function(gpu="L40S:1", @@ -480,10 +429,8 @@ def seed_grad_norm_references(): workstation — see the module docstring of ``grad_norm_regression.py`` for the local command and the ``_DEVICE_MAPPINGS`` table. """ - run_test( - "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && " - "FASTVIDEO_FA4=0 FASTVIDEO_GRADNORM_UPDATE=1 pytest ./fastvideo/tests/train/methods -vs -rs" - ) + run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && " + "FASTVIDEO_FA4=0 FASTVIDEO_GRADNORM_UPDATE=1 pytest ./fastvideo/tests/train/methods -vs -rs") @app.function(gpu="L40S:1", @@ -509,14 +456,9 @@ def run_eval_tests(): install_command='uv pip install -e ".[test,eval-full]"') -@app.function(gpu="L40S:1", - image=image, - timeout=3600, - secrets=[hf_secret, ci_env_secret]) +@app.function(gpu="L40S:1", image=image, timeout=3600, secrets=[hf_secret, ci_env_secret]) def run_lora_extraction_tests(): - run_test( - "hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/lora_extraction/test_lora_extraction.py" - ) + run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/lora_extraction/test_lora_extraction.py") @app.function(gpu="L40S:2", diff --git a/fastvideo/tests/modal/ssim_test.py b/fastvideo/tests/modal/ssim_test.py index 91033dea2..997637956 100644 --- a/fastvideo/tests/modal/ssim_test.py +++ b/fastvideo/tests/modal/ssim_test.py @@ -25,6 +25,7 @@ except ModuleNotFoundError: def resolve_uv_torch_backend(image_tag: str) -> str | None: return os.environ.get("UV_TORCH_BACKEND") + app = modal.App() model_vol = modal.Volume.from_name("hf-model-weights") @@ -39,8 +40,7 @@ print(f"Using image: {image_ref}") uv_torch_backend_override = resolve_uv_torch_backend(image_tag) image = ( - modal.Image.from_registry(image_ref, add_python="3.12") - .apt_install( + modal.Image.from_registry(image_ref, add_python="3.12").apt_install( "cmake", "pkg-config", "build-essential", @@ -52,25 +52,25 @@ image = ( "libsm6", "libxext6", "libxrender1", - ) - .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable") - .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc") - .env( - { - # INVARIANT: keep this image identical for every CI job at a given - # base digest -- per-job/per-commit values (BUILDKITE_*) must never - # be baked here as image layers or every job rebuilds its own image - # variant. repo/commit/PR already reach the container as - # run_ssim_partition() arguments. - "PATH": "/root/.cargo/bin:$PATH", - "IMAGE_VERSION": image_version, - **({"UV_TORCH_BACKEND": uv_torch_backend_override} if uv_torch_backend_override else {}), - # FA4 is opt-in (FASTVIDEO_FA4); the SSIM references were seeded - # with FA4 inference, so keep it enabled in CI. Caller override wins. - "FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1"), - } - ) -) + ).run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable") + .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc").env({ + # INVARIANT: keep this image identical for every CI job at a given + # base digest -- per-job/per-commit values (BUILDKITE_*) must never + # be baked here as image layers or every job rebuilds its own image + # variant. repo/commit/PR already reach the container as + # run_ssim_partition() arguments. + "PATH": + "/root/.cargo/bin:$PATH", + "IMAGE_VERSION": + image_version, + **({ + "UV_TORCH_BACKEND": uv_torch_backend_override + } if uv_torch_backend_override else {}), + # FA4 is opt-in (FASTVIDEO_FA4); the SSIM references were seeded + # with FA4 inference, so keep it enabled in CI. Caller override wins. + "FASTVIDEO_FA4": + os.environ.get("FASTVIDEO_FA4", "1"), + })) SSIM_NUM_GPUS = 4 SSIM_TERMINATE_TIMEOUT_S = 30 @@ -159,9 +159,9 @@ def _run_git_command(args: list[str]) -> str: def _normalize_git_repo_url(git_repo: str) -> str: if git_repo.startswith("git@github.com:"): - return "https://github.com/" + git_repo[len("git@github.com:") :] + return "https://github.com/" + git_repo[len("git@github.com:"):] if git_repo.startswith("ssh://git@github.com/"): - return "https://github.com/" + git_repo[len("ssh://git@github.com/") :] + return "https://github.com/" + git_repo[len("ssh://git@github.com/"):] return git_repo @@ -204,9 +204,7 @@ def _resolve_pull_request(pr_number: str) -> str: return "false" -def _resolve_hf_api_key( - hf_api_key: str, -) -> str: +def _resolve_hf_api_key(hf_api_key: str, ) -> str: if hf_api_key.strip(): return hf_api_key.strip() @@ -216,8 +214,7 @@ def _resolve_hf_api_key( return value raise RuntimeError( - "Hugging Face token is required. Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN; or pass --hf-api-key." - ) + "Hugging Face token is required. Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN; or pass --hf-api-key.") def _sanitize_path_fragment(value: str) -> str: @@ -268,13 +265,11 @@ def _print_local_reference_copy_command(quality_tier: str) -> None: _build_local_generated_download_dir(quality_tier), MODAL_DEVICE_REFERENCE_FOLDER, ) - print( - "To update local references from the downloaded Modal outputs, run:\n" - " python fastvideo/tests/ssim/reference_videos_cli.py copy-local " - f"--quality-tier {quality_tier} " - f"--generated-dir {generated_dir} " - f"--device-folder {MODAL_DEVICE_REFERENCE_FOLDER}" - ) + print("To update local references from the downloaded Modal outputs, run:\n" + " python fastvideo/tests/ssim/reference_videos_cli.py copy-local " + f"--quality-tier {quality_tier} " + f"--generated-dir {generated_dir} " + f"--device-folder {MODAL_DEVICE_REFERENCE_FOLDER}") def _count_video_files(root: str) -> int: @@ -318,11 +313,9 @@ def _sync_generated_videos_to_volume( num_videos = _count_video_files(absolute_dst) print(f"Raw generated videos exported to Modal volume path: {relative_dst} ({num_videos} video files).") - print( - "Download command:\n" - f" modal volume get hf-model-weights {relative_dst} " - f"{_build_local_generated_download_dir(quality_tier)}" - ) + print("Download command:\n" + f" modal volume get hf-model-weights {relative_dst} " + f"{_build_local_generated_download_dir(quality_tier)}") _print_local_reference_copy_command(quality_tier) return relative_dst @@ -395,11 +388,8 @@ def _discover_ssim_tasks( for filepath in test_files: file_name = os.path.basename(filepath) rel_path = f"./fastvideo/tests/ssim/{file_name}" - if selected_test_files and ( - filepath not in selected_test_files - and rel_path not in selected_test_files - and file_name not in selected_test_file_names - ): + if selected_test_files and (filepath not in selected_test_files and rel_path not in selected_test_files + and file_name not in selected_test_file_names): continue matched_test_files.add(file_name) @@ -418,19 +408,16 @@ def _discover_ssim_tasks( test_file=rel_path, required_gpus=required_gpus, model_id=model_id, - ) - ) + )) task_id += 1 else: if selected_model_ids: continue - tasks.append( - SSIMTask( - task_id=task_id, - test_file=rel_path, - required_gpus=required_gpus, - ) - ) + tasks.append(SSIMTask( + task_id=task_id, + test_file=rel_path, + required_gpus=required_gpus, + )) task_id += 1 if selected_test_files: @@ -711,12 +698,10 @@ def _schedule_ssim_tasks( max_required_gpus = max((task.required_gpus for task in pending_tasks), default=0) if max_required_gpus > len(available_gpu_ids): cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "") - raise RuntimeError( - "SSIM task requires " - f"{max_required_gpus} GPUs but only " - f"{len(available_gpu_ids)} are visible via " - f"CUDA_VISIBLE_DEVICES={cuda_visible_devices!r}." - ) + raise RuntimeError("SSIM task requires " + f"{max_required_gpus} GPUs but only " + f"{len(available_gpu_ids)} are visible via " + f"CUDA_VISIBLE_DEVICES={cuda_visible_devices!r}.") results = {} fail_fast_triggered = False log_dir = tempfile.mkdtemp(prefix="fastvideo-ssim-logs-") @@ -732,8 +717,8 @@ def _schedule_ssim_tasks( break task = pending_tasks.pop(next_task_index) - assigned_gpu_ids = available_gpu_ids[: task.required_gpus] - del available_gpu_ids[: task.required_gpus] + assigned_gpu_ids = available_gpu_ids[:task.required_gpus] + del available_gpu_ids[:task.required_gpus] running_task = _spawn_ssim_task( task=task, repo_root=repo_root, @@ -812,14 +797,11 @@ def _collect_task_summaries( status=result.status, returncode=result.returncode, log_content=log_content, - ) - ) + )) return summaries -def _print_combined_results( - partition_results: list[_PartitionResult | None], -) -> int: +def _print_combined_results(partition_results: list[_PartitionResult | None], ) -> int: """Print a unified report across all partitions.""" passed = [] failed = [] @@ -987,10 +969,8 @@ def run_ssim_tests( if pytest_k.strip(): print(f"Using pytest -k filter: {pytest_k}") if bootstrap_mode: - print( - "SSIM bootstrap mode enabled: missing references will upload " - "draft artifacts and xfail." - ) + print("SSIM bootstrap mode enabled: missing references will upload " + "draft artifacts and xfail.") quality_tier = _resolve_output_quality_tier(full_quality) if sync_generated_to_volume: resolved_subdir = _resolve_generated_volume_subdir( @@ -1001,10 +981,8 @@ def run_ssim_tests( generated_volume_subdir=resolved_subdir, quality_tier=quality_tier, ) - print( - "Raw generated videos will be saved to Modal volume path: " - f"{generated_volume_path}" - ) + print("Raw generated videos will be saved to Modal volume path: " + f"{generated_volume_path}") else: resolved_subdir = "" @@ -1029,8 +1007,7 @@ def run_ssim_tests( partition_index=i, num_partitions=NUM_PARTITIONS, **common_kwargs, - ) - for i in range(NUM_PARTITIONS) + ) for i in range(NUM_PARTITIONS) ] results: list[_PartitionResult | None] = [f.get() for f in futures] @@ -1041,11 +1018,9 @@ def run_ssim_tests( quality_tier=quality_tier, ) local_download_dir = _build_local_generated_download_dir(quality_tier) - print( - "To download raw generated videos locally, run:\n" - f" modal volume get hf-model-weights {download_src} " - f"{local_download_dir}" - ) + print("To download raw generated videos locally, run:\n" + f" modal volume get hf-model-weights {download_src} " + f"{local_download_dir}") _print_local_reference_copy_command(quality_tier) if exit_code != 0: sys.exit(exit_code) diff --git a/fastvideo/tests/modal/test_kernel_build_cache.py b/fastvideo/tests/modal/test_kernel_build_cache.py index 5047eb575..e043564e7 100644 --- a/fastvideo/tests/modal/test_kernel_build_cache.py +++ b/fastvideo/tests/modal/test_kernel_build_cache.py @@ -88,14 +88,14 @@ def _write_prebuilt_artifact(prebuilt_root: Path, name: str, cache_key: str) -> def _patch_stable_metadata(monkeypatch) -> None: for name in ( - "GPU_BACKEND", - "CMAKE_ARGS", - "CFLAGS", - "CXXFLAGS", - "LDFLAGS", - "TORCH_CUDA_ARCH_LIST", - "FASTVIDEO_CONTAINER_IMAGE_REF", - "CUDACXX", + "GPU_BACKEND", + "CMAKE_ARGS", + "CFLAGS", + "CXXFLAGS", + "LDFLAGS", + "TORCH_CUDA_ARCH_LIST", + "FASTVIDEO_CONTAINER_IMAGE_REF", + "CUDACXX", ): monkeypatch.delenv(name, raising=False) monkeypatch.setattr(kernel_build_cache, "_kernel_source_hash", lambda repo_root: {"source_hash": "source"}) @@ -115,8 +115,15 @@ def _patch_stable_metadata(monkeypatch) -> None: kernel_build_cache, "_compiler_libc_metadata", lambda: { - "compiler": {"cc": {"path": "/usr/bin/gcc", "version": "gcc 11.4.0"}}, - "libc": {"ldd_version": "ldd 2.35"}, + "compiler": { + "cc": { + "path": "/usr/bin/gcc", + "version": "gcc 11.4.0" + } + }, + "libc": { + "ldd_version": "ldd 2.35" + }, }, ) monkeypatch.setattr( @@ -188,8 +195,15 @@ def test_cache_key_changes_with_compiler_or_torch_abi(monkeypatch, tmp_path) -> kernel_build_cache, "_compiler_libc_metadata", lambda: { - "compiler": {"cc": {"path": "/opt/gcc-12/bin/gcc", "version": "gcc 12.3.0"}}, - "libc": {"ldd_version": "ldd 2.35"}, + "compiler": { + "cc": { + "path": "/opt/gcc-12/bin/gcc", + "version": "gcc 12.3.0" + } + }, + "libc": { + "ldd_version": "ldd 2.35" + }, }, ) compiler_changed = kernel_build_cache._build_metadata(tmp_path) diff --git a/fastvideo/tests/modal/test_pr_test.py b/fastvideo/tests/modal/test_pr_test.py index 7190d0bd0..c55a7a695 100644 --- a/fastvideo/tests/modal/test_pr_test.py +++ b/fastvideo/tests/modal/test_pr_test.py @@ -44,6 +44,7 @@ class _FakeSecret: class _FakeApp: def function(self, *_args, **_kwargs): + def decorator(func): return func @@ -64,8 +65,7 @@ def _load_pr_test_module(monkeypatch): monkeypatch.setitem(sys.modules, "modal", fake_modal) monkeypatch.setitem(sys.modules, "modal_image_utils", fake_image_utils) module_path = Path(__file__).with_name("pr_test.py") - spec = importlib.util.spec_from_file_location( - "modal_pr_test_under_test", module_path) + spec = importlib.util.spec_from_file_location("modal_pr_test_under_test", module_path) module = importlib.util.module_from_spec(spec) assert spec.loader is not None spec.loader.exec_module(module) @@ -87,8 +87,7 @@ def test_checkout_repository_retries_clone_and_fetches_pr_ref(monkeypatch): monkeypatch.setattr( module.shutil, "rmtree", - lambda path, ignore_errors: cleanup_paths.append( - (path, ignore_errors)), + lambda path, ignore_errors: cleanup_paths.append((path, ignore_errors)), ) monkeypatch.setattr(module.time, "sleep", sleep_seconds.append) @@ -121,11 +120,26 @@ def test_checkout_repository_retries_clone_and_fetches_pr_ref(monkeypatch): ] assert sleep_seconds == [5] assert [kwargs for _, kwargs in commands] == [ - {"cwd": "/", "check": False}, - {"cwd": "/", "check": False}, - {"cwd": "/tmp/FastVideo", "check": False}, - {"cwd": "/tmp/FastVideo", "check": False}, - {"cwd": "/tmp/FastVideo", "check": False}, + { + "cwd": "/", + "check": False + }, + { + "cwd": "/", + "check": False + }, + { + "cwd": "/tmp/FastVideo", + "check": False + }, + { + "cwd": "/tmp/FastVideo", + "check": False + }, + { + "cwd": "/tmp/FastVideo", + "check": False + }, ] fetch_command = commands[2][0] @@ -187,8 +201,7 @@ def test_git_retry_exhaustion_is_bounded_and_cleans_each_attempt(monkeypatch): monkeypatch.setattr( module.shutil, "rmtree", - lambda path, ignore_errors: cleanup_paths.append( - (path, ignore_errors)), + lambda path, ignore_errors: cleanup_paths.append((path, ignore_errors)), ) monkeypatch.setattr(module.time, "sleep", sleep_seconds.append) @@ -215,14 +228,12 @@ def test_git_retry_exhaustion_is_bounded_and_cleans_each_attempt(monkeypatch): ("https://example.com/repo.git", "0123456789abcdef", "0"), ], ) -def test_checkout_repository_rejects_invalid_buildkite_values( - monkeypatch, git_repo, git_commit, pr_number): +def test_checkout_repository_rejects_invalid_buildkite_values(monkeypatch, git_repo, git_commit, pr_number): module = _load_pr_test_module(monkeypatch) monkeypatch.setattr( module.subprocess, "run", - lambda *_args, **_kwargs: pytest.fail( - "git must not run for invalid input"), + lambda *_args, **_kwargs: pytest.fail("git must not run for invalid input"), ) with pytest.raises(RuntimeError): @@ -236,16 +247,13 @@ def test_checkout_repository_rejects_invalid_buildkite_values( (False, ""), ], ) -def test_run_test_command_composes_valid_post_checkout_shell( - monkeypatch, build_kernel, install_command): +def test_run_test_command_composes_valid_post_checkout_shell(monkeypatch, build_kernel, install_command): module = _load_pr_test_module(monkeypatch) real_run = subprocess.run events = [] - monkeypatch.setenv( - "BUILDKITE_REPO", "https://github.com/hao-ai-lab/FastVideo.git") - monkeypatch.setenv( - "BUILDKITE_COMMIT", "0123456789abcdef0123456789abcdef01234567") + monkeypatch.setenv("BUILDKITE_REPO", "https://github.com/hao-ai-lab/FastVideo.git") + monkeypatch.setenv("BUILDKITE_COMMIT", "0123456789abcdef0123456789abcdef01234567") monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "false") monkeypatch.setattr( module, @@ -293,8 +301,7 @@ def test_run_test_command_uses_nonshared_kernel_install_before_tests(monkeypatch monkeypatch.setenv("BUILDKITE_REPO", "https://example.com/FastVideo.git") monkeypatch.setenv("BUILDKITE_COMMIT", "0123456789abcdef") monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "false") - monkeypatch.setattr(module, "_checkout_repository", - lambda *_args: None) + monkeypatch.setattr(module, "_checkout_repository", lambda *_args: None) monkeypatch.setattr(subprocess, "run", fake_run) module.run_test_command("pytest fastvideo/tests/api -q", build_kernel=True) @@ -318,8 +325,8 @@ def test_run_unit_test_collects_modal_cache_runner_tests(monkeypatch): assert len(commands) == 1 for test_path in ( - "./fastvideo/tests/modal/test_kernel_build_cache.py", - "./fastvideo/tests/modal/test_pr_test.py", - "./fastvideo/tests/modal/test_ssim_test.py", + "./fastvideo/tests/modal/test_kernel_build_cache.py", + "./fastvideo/tests/modal/test_pr_test.py", + "./fastvideo/tests/modal/test_ssim_test.py", ): assert test_path in commands[0] diff --git a/fastvideo/tests/modal/test_ssim_test.py b/fastvideo/tests/modal/test_ssim_test.py index 574307740..5aa46323c 100644 --- a/fastvideo/tests/modal/test_ssim_test.py +++ b/fastvideo/tests/modal/test_ssim_test.py @@ -34,12 +34,14 @@ class _FakeVolume: class _FakeApp: def function(self, *_args, **_kwargs): + def decorator(func): return func return decorator def local_entrypoint(self, *_args, **_kwargs): + def decorator(func): return func diff --git a/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py b/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py index 1adf817f3..b851d6de5 100644 --- a/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py +++ b/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py @@ -20,7 +20,6 @@ PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/pipelines/preprocess/v1_preprocess.py LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "crush-smol_processed_t2v")) - # training NUM_GPUS_PER_NODE_TRAINING = "4" TRAINING_ENTRY_FILE_PATH = "fastvideo/training/wan_distillation_pipeline.py" @@ -28,6 +27,7 @@ LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "combined_pa LOCAL_VALIDATION_DATASET_FILE = "examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs")) + def download_data(): # create the data dir if it doesn't exist data_dir = Path(DATA_DIR) @@ -45,17 +45,17 @@ def download_data(): token=os.environ.get("HF_TOKEN"), # In case authentication is needed ) print(f"Download completed successfully. Files downloaded to: {result}") - + # Verify the download if not LOCAL_RAW_DATA_DIR.exists(): raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist") - + # List downloaded files print("Downloaded files:") for file in LOCAL_RAW_DATA_DIR.rglob("*"): if file.is_file(): print(f" - {file.relative_to(LOCAL_RAW_DATA_DIR)}") - + except Exception as e: print(f"Error during download: {str(e)}") raise @@ -70,23 +70,39 @@ def run_preprocessing(): # Run torchrun command cmd = [ "torchrun", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING, + "--nnodes", + NUM_NODES, + "--nproc_per_node", + NUM_GPUS_PER_NODE_PREPROCESSING, PREPROCESSING_ENTRY_FILE_PATH, - "--model_path", MODEL_PATH, - "--seed", "42", - "--data_merge_path", os.path.join(LOCAL_RAW_DATA_DIR, "merge.txt"), - "--preprocess_video_batch_size", "1", - "--max_height", "480", - "--max_width", "832", - "--num_frames", "81", - "--dataloader_num_workers", "0", - "--output_dir", LOCAL_PREPROCESSED_DATA_DIR, - "--train_fps", "16", - "--samples_per_file", "1", - "--flush_frequency", "1", - "--video_length_tolerance_range", "5", - "--preprocess_task", "t2v", + "--model_path", + MODEL_PATH, + "--seed", + "42", + "--data_merge_path", + os.path.join(LOCAL_RAW_DATA_DIR, "merge.txt"), + "--preprocess_video_batch_size", + "1", + "--max_height", + "480", + "--max_width", + "832", + "--num_frames", + "81", + "--dataloader_num_workers", + "0", + "--output_dir", + LOCAL_PREPROCESSED_DATA_DIR, + "--train_fps", + "16", + "--samples_per_file", + "1", + "--flush_frequency", + "1", + "--video_length_tolerance_range", + "5", + "--preprocess_task", + "t2v", ] process = subprocess.run(cmd, check=True) @@ -95,53 +111,98 @@ def run_preprocessing(): def run_training(): cmd = [ "torchrun", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING, + "--nnodes", + NUM_NODES, + "--nproc_per_node", + NUM_GPUS_PER_NODE_TRAINING, TRAINING_ENTRY_FILE_PATH, - "--model_path", MODEL_PATH, - "--inference_mode", "False", - "--pretrained_model_name_or_path", MODEL_PATH, - "--data_path", LOCAL_TRAINING_DATA_DIR, - "--validation_dataset_file", LOCAL_VALIDATION_DATASET_FILE, - "--train_batch_size", "1", - "--num_latent_t", "8", - "--num_gpus", NUM_GPUS_PER_NODE_TRAINING, - "--sp_size", "1", - "--tp_size", "1", - "--hsdp_replicate_dim", "1", - "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING, - "--train_sp_batch_size", "1", - "--dataloader_num_workers", "10", - "--gradient_accumulation_steps", "1", - "--max_train_steps", "501", - "--learning_rate", "2e-6", - "--fake_score_learning_rate", "2e-6", - "--mixed_precision", "bf16", - "--training_state_checkpointing_steps", "1000", - "--weight_only_checkpointing_steps", "1000", - "--validation_steps", "50", - "--validation_sampling_steps", "3", + "--model_path", + MODEL_PATH, + "--inference_mode", + "False", + "--pretrained_model_name_or_path", + MODEL_PATH, + "--data_path", + LOCAL_TRAINING_DATA_DIR, + "--validation_dataset_file", + LOCAL_VALIDATION_DATASET_FILE, + "--train_batch_size", + "1", + "--num_latent_t", + "8", + "--num_gpus", + NUM_GPUS_PER_NODE_TRAINING, + "--sp_size", + "1", + "--tp_size", + "1", + "--hsdp_replicate_dim", + "1", + "--hsdp_shard_dim", + NUM_GPUS_PER_NODE_TRAINING, + "--train_sp_batch_size", + "1", + "--dataloader_num_workers", + "10", + "--gradient_accumulation_steps", + "1", + "--max_train_steps", + "501", + "--learning_rate", + "2e-6", + "--fake_score_learning_rate", + "2e-6", + "--mixed_precision", + "bf16", + "--training_state_checkpointing_steps", + "1000", + "--weight_only_checkpointing_steps", + "1000", + "--validation_steps", + "50", + "--validation_sampling_steps", + "3", "--log_validation", - "--checkpoints_total_limit", "3", - "--ema_start_step", "0", - "--training_cfg_rate", "0.0", - "--output_dir", LOCAL_OUTPUT_DIR, - "--tracker_project_name", "ci_wan_t2v_dmd_overfit", - "--num_height", "480", - "--num_width", "832", - "--num_frames", "81", - "--flow_shift", "8", - "--validation_guidance_scale", "6.0", - "--weight_decay", "0.01", - "--generator_update_interval", "5", - "--dmd_denoising_steps", "1000,757,522", - "--min_timestep_ratio", "0.02", - "--max_timestep_ratio", "0.98", - "--seed", "1000", - "--real_score_guidance_scale", "3.5", - "--dit_precision", "fp32", - "--max_grad_norm", "1.0", - "--enable_gradient_checkpointing_type", "full", + "--checkpoints_total_limit", + "3", + "--ema_start_step", + "0", + "--training_cfg_rate", + "0.0", + "--output_dir", + LOCAL_OUTPUT_DIR, + "--tracker_project_name", + "ci_wan_t2v_dmd_overfit", + "--num_height", + "480", + "--num_width", + "832", + "--num_frames", + "81", + "--flow_shift", + "8", + "--validation_guidance_scale", + "6.0", + "--weight_decay", + "0.01", + "--generator_update_interval", + "5", + "--dmd_denoising_steps", + "1000,757,522", + "--min_timestep_ratio", + "0.02", + "--max_timestep_ratio", + "0.98", + "--seed", + "1000", + "--real_score_guidance_scale", + "3.5", + "--dit_precision", + "fp32", + "--max_grad_norm", + "1.0", + "--enable_gradient_checkpointing_type", + "full", ] print(f"Running training with command: {cmd}") @@ -160,7 +221,6 @@ def test_e2e_overfit_single_sample(): final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4") print(f"final_validation_video_file: {final_validation_video_file}") - # Ensure both files exist assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}" assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}" diff --git a/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py b/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py index 88c409f40..4f82a6fd9 100644 --- a/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py +++ b/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py @@ -20,7 +20,6 @@ PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/pipelines/preprocess/v1_preprocess.py LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "cats_preprocessed_data_i2v")) - # training NUM_GPUS_PER_NODE_TRAINING = "4" TRAINING_ENTRY_FILE_PATH = "fastvideo/training/wan_i2v_training_pipeline.py" @@ -28,6 +27,7 @@ LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "combined_pa LOCAL_VALIDATION_DATASET_FILE = os.path.join(LOCAL_RAW_DATA_DIR, "validation_i2v_prompt_1_sample.json") LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs")) + def download_data(): # create the data dir if it doesn't exist data_dir = Path(DATA_DIR) @@ -45,17 +45,17 @@ def download_data(): token=os.environ.get("HF_TOKEN"), # In case authentication is needed ) print(f"Download completed successfully. Files downloaded to: {result}") - + # Verify the download if not LOCAL_RAW_DATA_DIR.exists(): raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist") - + # List downloaded files print("Downloaded files:") for file in LOCAL_RAW_DATA_DIR.rglob("*"): if file.is_file(): print(f" - {file.relative_to(LOCAL_RAW_DATA_DIR)}") - + except Exception as e: print(f"Error during download: {str(e)}") raise @@ -65,22 +65,37 @@ def run_preprocessing(): # Run torchrun command cmd = [ "torchrun", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING, + "--nnodes", + NUM_NODES, + "--nproc_per_node", + NUM_GPUS_PER_NODE_PREPROCESSING, PREPROCESSING_ENTRY_FILE_PATH, - "--model_path", MODEL_PATH, - "--data_merge_path", os.path.join(LOCAL_RAW_DATA_DIR, "merge_1_sample.txt"), - "--preprocess_video_batch_size", "1", - "--max_height", "480", - "--max_width", "832", - "--num_frames", "77", - "--dataloader_num_workers", "0", - "--output_dir", LOCAL_PREPROCESSED_DATA_DIR, - "--train_fps", "16", - "--samples_per_file", "1", - "--flush_frequency", "1", - "--video_length_tolerance_range", "5", - "--preprocess_task", "i2v", + "--model_path", + MODEL_PATH, + "--data_merge_path", + os.path.join(LOCAL_RAW_DATA_DIR, "merge_1_sample.txt"), + "--preprocess_video_batch_size", + "1", + "--max_height", + "480", + "--max_width", + "832", + "--num_frames", + "77", + "--dataloader_num_workers", + "0", + "--output_dir", + LOCAL_PREPROCESSED_DATA_DIR, + "--train_fps", + "16", + "--samples_per_file", + "1", + "--flush_frequency", + "1", + "--video_length_tolerance_range", + "5", + "--preprocess_task", + "i2v", ] process = subprocess.run(cmd, check=True) @@ -89,48 +104,87 @@ def run_preprocessing(): def run_training(): cmd = [ "torchrun", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING, + "--nnodes", + NUM_NODES, + "--nproc_per_node", + NUM_GPUS_PER_NODE_TRAINING, TRAINING_ENTRY_FILE_PATH, - "--model_path", MODEL_PATH, - "--inference_mode", "False", - "--pretrained_model_name_or_path", MODEL_PATH, - "--data_path", LOCAL_TRAINING_DATA_DIR, - "--validation_dataset_file", LOCAL_VALIDATION_DATASET_FILE, - "--train_batch_size", "1", - "--num_latent_t", "8", - "--num_gpus", NUM_GPUS_PER_NODE_TRAINING, - "--sp_size", NUM_GPUS_PER_NODE_TRAINING, - "--tp_size", "1", - "--hsdp_replicate_dim", "1", - "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING, - "--num_gpus", NUM_GPUS_PER_NODE_TRAINING, - "--train_sp_batch_size", "1", - "--dataloader_num_workers", "10", - "--gradient_accumulation_steps", "1", - "--max_train_steps", "901", - "--learning_rate", "1e-5", - "--mixed_precision", "bf16", - "--weight_only_checkpointing_steps", "6000", - "--training_state_checkpointing_steps", "6000", - "--validation_steps", "100", - "--validation_sampling_steps", "50", + "--model_path", + MODEL_PATH, + "--inference_mode", + "False", + "--pretrained_model_name_or_path", + MODEL_PATH, + "--data_path", + LOCAL_TRAINING_DATA_DIR, + "--validation_dataset_file", + LOCAL_VALIDATION_DATASET_FILE, + "--train_batch_size", + "1", + "--num_latent_t", + "8", + "--num_gpus", + NUM_GPUS_PER_NODE_TRAINING, + "--sp_size", + NUM_GPUS_PER_NODE_TRAINING, + "--tp_size", + "1", + "--hsdp_replicate_dim", + "1", + "--hsdp_shard_dim", + NUM_GPUS_PER_NODE_TRAINING, + "--num_gpus", + NUM_GPUS_PER_NODE_TRAINING, + "--train_sp_batch_size", + "1", + "--dataloader_num_workers", + "10", + "--gradient_accumulation_steps", + "1", + "--max_train_steps", + "901", + "--learning_rate", + "1e-5", + "--mixed_precision", + "bf16", + "--weight_only_checkpointing_steps", + "6000", + "--training_state_checkpointing_steps", + "6000", + "--validation_steps", + "100", + "--validation_sampling_steps", + "50", "--log_validation", - "--checkpoints_total_limit", "3", - "--ema_start_step", "0", - "--training_cfg_rate", "0.1", - "--output_dir", LOCAL_OUTPUT_DIR, - "--tracker_project_name", "wan_i2v_finetune_overfit_ci", - "--num_height", "480", - "--num_width", "832", - "--num_frames", "81", - "--validation_guidance_scale", "6.0", - "--num_euler_timesteps", "50", - "--multi_phased_distill_schedule", "4000-1", - "--weight_decay", "0.01", + "--checkpoints_total_limit", + "3", + "--ema_start_step", + "0", + "--training_cfg_rate", + "0.1", + "--output_dir", + LOCAL_OUTPUT_DIR, + "--tracker_project_name", + "wan_i2v_finetune_overfit_ci", + "--num_height", + "480", + "--num_width", + "832", + "--num_frames", + "81", + "--validation_guidance_scale", + "6.0", + "--num_euler_timesteps", + "50", + "--multi_phased_distill_schedule", + "4000-1", + "--weight_decay", + "0.01", "--not_apply_cfg_solver", - "--dit_precision", "fp32", - "--max_grad_norm", "1.0", + "--dit_precision", + "fp32", + "--max_grad_norm", + "1.0", ] print(f"Running training with command: {cmd}") @@ -149,7 +203,6 @@ def test_e2e_overfit_single_sample(): final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4") print(f"final_validation_video_file: {final_validation_video_file}") - # Ensure both files exist assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}" assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}" diff --git a/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py b/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py index be948e77d..43971b7fe 100644 --- a/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py +++ b/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py @@ -102,8 +102,7 @@ STAGE1_DIFFUSERS_DIR = E2E_ROOT / "stage1_diffusers" STAGE2_DIFFUSERS_DIR = E2E_ROOT / "stage2_diffusers" GENERATED_VIDEO_DIR = E2E_ROOT / "generated" -STAGE1_CONFIG = (REPO_ROOT / "examples" / "train" / "configs" / "fine_tuning" / "kandinsky5" / - "t2v_480p_qat.yaml") +STAGE1_CONFIG = (REPO_ROOT / "examples" / "train" / "configs" / "fine_tuning" / "kandinsky5" / "t2v_480p_qat.yaml") STAGE2_CONFIG = (REPO_ROOT / "examples" / "train" / "configs" / "distribution_matching" / "kandinsky5" / "dmd2_t2v_480p_qat.yaml") @@ -194,9 +193,8 @@ def _synthesize_single_sample() -> None: (768, 512), ) if not writer.isOpened(): - raise RuntimeError( - f"cv2.VideoWriter failed to open {video_path} (mp4v codec unavailable?) -- " - "the fixture clip would be empty and preprocessing would fail on it.") + raise RuntimeError(f"cv2.VideoWriter failed to open {video_path} (mp4v codec unavailable?) -- " + "the fixture clip would be empty and preprocessing would fail on it.") rng = np.random.default_rng(0) for _ in range(121): frame = rng.integers(0, 256, size=(512, 768, 3), dtype=np.uint8) @@ -228,7 +226,8 @@ def _run_preprocessing() -> None: env["KANDINSKY5_OVERFIT_DATA_DIR"] = str(RAW_DATA_DIR) env["KANDINSKY5_OVERFIT_OUTPUT_DIR"] = str(PREPROCESSED_DIR) cmd = [ - sys.executable, "-m", + sys.executable, + "-m", "fastvideo.pipelines.preprocess.preprocess_kandinsky5_overfit", ] subprocess.run(cmd, cwd=str(REPO_ROOT), env=env, check=True) @@ -247,11 +246,15 @@ def _export_dcp_to_diffusers(checkpoint_dir: Path, output_dir: Path) -> None: inside the next launch that loads the directory. """ cmd = [ - sys.executable, "-m", + sys.executable, + "-m", "fastvideo.train.entrypoint.dcp_to_diffusers", - "--checkpoint", str(checkpoint_dir), - "--output-dir", str(output_dir), - "--role", "student", + "--checkpoint", + str(checkpoint_dir), + "--output-dir", + str(output_dir), + "--role", + "student", "--overwrite", "--verify", ] @@ -265,27 +268,42 @@ def _latest_checkpoint(output_dir: Path) -> Path: def _run_stage(config: Path, output_dir: Path, *, max_train_steps: int, extra_overrides: list[str], - env_overrides: dict[str, str]) -> None: + env_overrides: dict[str, str]) -> None: cmd = [ - sys.executable, "-m", "torch.distributed.run", - "--nnodes", "1", - "--nproc_per_node", NUM_GPUS, - "-m", "fastvideo.train.entrypoint.train", - "--config", str(config), - "--training.data.data_path", str(PREPROCESSED_DIR), - "--training.distributed.num_gpus", NUM_GPUS, - "--training.distributed.hsdp_shard_dim", NUM_GPUS, - "--training.loop.max_train_steps", str(max_train_steps), - "--training.checkpoint.output_dir", str(output_dir), - "--training.checkpoint.training_state_checkpointing_steps", str(max_train_steps), + sys.executable, + "-m", + "torch.distributed.run", + "--nnodes", + "1", + "--nproc_per_node", + NUM_GPUS, + "-m", + "fastvideo.train.entrypoint.train", + "--config", + str(config), + "--training.data.data_path", + str(PREPROCESSED_DIR), + "--training.distributed.num_gpus", + NUM_GPUS, + "--training.distributed.hsdp_shard_dim", + NUM_GPUS, + "--training.loop.max_train_steps", + str(max_train_steps), + "--training.checkpoint.output_dir", + str(output_dir), + "--training.checkpoint.training_state_checkpointing_steps", + str(max_train_steps), # Explicitly disable resume for BOTH stages, even though # _clean_previous_artifacts() already removed the output dirs: the # stage-1 YAML defaults to resume_from_checkpoint: latest, and this # test's correctness must not silently depend on that default (or a # partially-failed cleanup) ever changing. - "--training.checkpoint.resume_from_checkpoint", "none", - "--callbacks.validation.every_steps", str(max_train_steps), - "--callbacks.validation.dataset_file", str(PREPROCESSED_DIR / "validation_prompts.json"), + "--training.checkpoint.resume_from_checkpoint", + "none", + "--callbacks.validation.every_steps", + str(max_train_steps), + "--callbacks.validation.dataset_file", + str(PREPROCESSED_DIR / "validation_prompts.json"), *extra_overrides, ] env = dict(os.environ) @@ -316,15 +334,15 @@ def _generate_from_export(export_dir: Path, output_dir: Path) -> Path: env.pop("FASTVIDEO_ATTENTION_BACKEND", None) cmd = [ sys.executable, - str(THIS_FILE), "--generate", + str(THIS_FILE), + "--generate", str(export_dir), str(output_dir), ] subprocess.run(cmd, cwd=str(REPO_ROOT), env=env, check=True) videos = sorted(Path(output_dir).glob("**/*.mp4")) - assert len(videos) == 1, ( - f"expected exactly one generated video under {output_dir}, found {videos}") + assert len(videos) == 1, (f"expected exactly one generated video under {output_dir}, found {videos}") return videos[0] @@ -420,11 +438,10 @@ def _assert_matches_reference(video_path: Path) -> None: "review it by eye and commit it. This bootstrap run " "compares the video against itself (SSIM=1) below.") else: - pytest.fail( - f"reference video missing at {REFERENCE_VIDEO}. This test " - "must not pass on artifact existence alone -- run once on a " - f"sanctioned GPU box with {WRITE_REFERENCE_ENV}=1, review " - "the written video, and commit it (see module docstring).") + pytest.fail(f"reference video missing at {REFERENCE_VIDEO}. This test " + "must not pass on artifact existence alone -- run once on a " + f"sanctioned GPU box with {WRITE_REFERENCE_ENV}=1, review " + "the written video, and commit it (see module docstring).") from fastvideo.tests.utils import compute_video_ssim_torchvision @@ -446,7 +463,9 @@ def _assert_matches_reference(video_path: Path) -> None: @pytest.mark.nightly def test_e2e_kandinsky5_dmd_overfit_single_sample(): if not STAGE1_CONFIG.exists() or not STAGE2_CONFIG.exists(): - pytest.skip("Kandinsky5 QAT configs not found -- see examples/train/configs/{fine_tuning,distribution_matching}/kandinsky5/") + pytest.skip( + "Kandinsky5 QAT configs not found -- see examples/train/configs/{fine_tuning,distribution_matching}/kandinsky5/" + ) os.environ.setdefault("WANDB_MODE", "offline") @@ -469,7 +488,8 @@ def test_e2e_kandinsky5_dmd_overfit_single_sample(): # eye at bootstrap, and stable across runs for the SSIM oracle # (noise output decorrelates under any training # nondeterminism; structured output does not). - "--training.optimizer.learning_rate", "1e-6", + "--training.optimizer.learning_rate", + "1e-6", ], env_overrides={"FASTVIDEO_ATTENTION_BACKEND": "ATTN_QAT_TRAIN"}, ) @@ -485,9 +505,12 @@ def test_e2e_kandinsky5_dmd_overfit_single_sample(): STAGE2_OUTPUT_DIR, max_train_steps=3, extra_overrides=[ - "--models.student.init_from", str(stage1_diffusers_dir), - "--models.teacher.init_from", str(stage1_diffusers_dir), - "--models.critic.init_from", str(stage1_diffusers_dir), + "--models.student.init_from", + str(stage1_diffusers_dir), + "--models.teacher.init_from", + str(stage1_diffusers_dir), + "--models.critic.init_from", + str(stage1_diffusers_dir), # The recipe's generator_update_interval of 5 would mean ZERO # student updates in a 3-step run (the trainer iterates 1..3 # and DMD2Method updates the generator only when @@ -495,12 +518,12 @@ def test_e2e_kandinsky5_dmd_overfit_single_sample(): # would just be the stage-1 weights, and this test could not # catch a broken student backward/optimizer path. Update the # generator every step instead. - "--method.generator_update_interval", "1", + "--method.generator_update_interval", + "1", ], env_overrides={"FASTVIDEO_ATTENTION_BACKEND": "ATTN_QAT_TRAIN"}, ) - assert any(STAGE2_OUTPUT_DIR.glob("*.mp4")), ( - f"no stage-2 validation video produced under {STAGE2_OUTPUT_DIR}") + assert any(STAGE2_OUTPUT_DIR.glob("*.mp4")), (f"no stage-2 validation video produced under {STAGE2_OUTPUT_DIR}") # The actual deliverable of this recipe is the exported stage-2 # student: export it, strict-reload it (--verify), then instantiate it diff --git a/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py b/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py index 373490fdd..5ffa76e82 100644 --- a/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py +++ b/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py @@ -76,8 +76,7 @@ def download_data(): repo_type="dataset", token=os.environ.get("HF_TOKEN"), ) - assert LOCAL_RAW_DATA_DIR.exists(), ( - f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist") + assert LOCAL_RAW_DATA_DIR.exists(), (f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist") def run_preprocessing(case: dict): @@ -96,17 +95,22 @@ def run_training(case: dict): env.setdefault("WANDB_MODE", "offline") cmd = [ "torchrun", - "--nnodes", "1", - "--nproc_per_node", NUM_GPUS_TRAINING, - "-m", "fastvideo.train.entrypoint.train", - "--config", case["config"], - "--training.checkpoint.output_dir", str(case["out_dir"]), - "--training.data.data_path", str(case["prep_dir"]), + "--nnodes", + "1", + "--nproc_per_node", + NUM_GPUS_TRAINING, + "-m", + "fastvideo.train.entrypoint.train", + "--config", + case["config"], + "--training.checkpoint.output_dir", + str(case["out_dir"]), + "--training.data.data_path", + str(case["prep_dir"]), "--callbacks.validation.dataset_file", str(case["prep_dir"] / "validation_prompts.json"), ] - if (case.get("requires_nvfp4") - and torch.cuda.get_device_capability(0) != (12, 0)): + if (case.get("requires_nvfp4") and torch.cuda.get_device_capability(0) != (12, 0)): cmd.extend(["--callbacks.validation.attn_qat_infer", "false"]) print(f"Running training: {cmd}") subprocess.run(cmd, check=True, env=env) @@ -115,14 +119,11 @@ def run_training(case: dict): def test_run_training_disables_sm120_only_validation_on_gb200(monkeypatch): commands = [] monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _: (10, 0)) - monkeypatch.setattr(subprocess, "run", - lambda cmd, **_: commands.append(cmd)) + monkeypatch.setattr(subprocess, "run", lambda cmd, **_: commands.append(cmd)) run_training(_CASES["ltx2_nvfp4_qat"]) - assert commands[0][-2:] == [ - "--callbacks.validation.attn_qat_infer", "false" - ] + assert commands[0][-2:] == ["--callbacks.validation.attn_qat_infer", "false"] def _validation_videos_by_step(out_dir: Path) -> dict[int, str]: @@ -143,14 +144,11 @@ def _validation_videos_by_step(out_dir: Path) -> dict[int, str]: def test_e2e_ltx2_overfit_new_stack(case_id: str): case = _CASES[case_id] if case.get("requires_nvfp4"): - if (not torch.cuda.is_available() or - torch.cuda.get_device_capability(0) < (10, 0)): + if (not torch.cuda.is_available() or torch.cuda.get_device_capability(0) < (10, 0)): pytest.skip("NVFP4 QAT requires an SM100+ GPU") try: from flashinfer import ( # noqa: F401 - SfLayout, - mm_fp4, - nvfp4_quantize, + SfLayout, mm_fp4, nvfp4_quantize, ) except ImportError: pytest.skip("NVFP4 QAT requires flashinfer with FP4 kernels") @@ -159,8 +157,7 @@ def test_e2e_ltx2_overfit_new_stack(case_id: str): run_training(case) reference_video = case["prep_dir"] / "training_sample_0.mp4" - assert reference_video.exists(), ( - f"Reference (preprocessed training clip) not found at {reference_video}") + assert reference_video.exists(), (f"Reference (preprocessed training clip) not found at {reference_video}") videos = _validation_videos_by_step(case["out_dir"]) assert videos, f"No validation videos found under {case['out_dir']}" @@ -168,26 +165,24 @@ def test_e2e_ltx2_overfit_new_stack(case_id: str): final_video = videos[final_step] print(f"Final validation video (step {final_step}): {final_video}") - final_mean, final_min, final_max = compute_video_ssim_torchvision( - str(reference_video), final_video, use_ms_ssim=True) + final_mean, final_min, final_max = compute_video_ssim_torchvision(str(reference_video), + final_video, + use_ms_ssim=True) print(f"\n===== MS-SSIM vs training clip at step {final_step} ({case_id}) =====") print(f"Mean: {final_mean:.4f} Min: {final_min:.4f} Max: {final_max:.4f}") results = {"case": case_id, "final_step": final_step, "final_mean_ssim": final_mean} baseline_step = min(videos) if baseline_step != final_step: - base_mean, _, _ = compute_video_ssim_torchvision( - str(reference_video), videos[baseline_step], use_ms_ssim=True) + base_mean, _, _ = compute_video_ssim_torchvision(str(reference_video), videos[baseline_step], use_ms_ssim=True) print(f"Baseline (step {baseline_step}) mean MS-SSIM: {base_mean:.4f}") results["baseline_mean_ssim"] = base_mean - assert final_mean > base_mean, ( - f"Overfitting did not improve similarity to the training clip: " - f"step {final_step} mean SSIM {final_mean:.4f} <= " - f"step {baseline_step} mean SSIM {base_mean:.4f}") + assert final_mean > base_mean, (f"Overfitting did not improve similarity to the training clip: " + f"step {final_step} mean SSIM {final_mean:.4f} <= " + f"step {baseline_step} mean SSIM {base_mean:.4f}") print(json.dumps(results)) - assert final_mean > 0.5, ( - f"Mean MS-SSIM vs the training clip is below 0.5: {final_mean:.4f}") + assert final_mean > 0.5, (f"Mean MS-SSIM vs the training clip is below 0.5: {final_mean:.4f}") if __name__ == "__main__": diff --git a/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py b/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py index 93f44eeb3..1d1f8f9c7 100644 --- a/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py +++ b/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py @@ -20,7 +20,6 @@ PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/pipelines/preprocess/v1_preprocess.py LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "cats_preprocessed_data")) - # training NUM_GPUS_PER_NODE_TRAINING = "4" TRAINING_ENTRY_FILE_PATH = "fastvideo/training/wan_training_pipeline.py" @@ -28,6 +27,7 @@ LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "combined_pa LOCAL_VALIDATION_DATASET_FILE = os.path.join(LOCAL_RAW_DATA_DIR, "validation_prompt_1_sample.json") LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs")) + def download_data(): # create the data dir if it doesn't exist data_dir = Path(DATA_DIR) @@ -45,17 +45,17 @@ def download_data(): token=os.environ.get("HF_TOKEN"), # In case authentication is needed ) print(f"Download completed successfully. Files downloaded to: {result}") - + # Verify the download if not LOCAL_RAW_DATA_DIR.exists(): raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist") - + # List downloaded files print("Downloaded files:") for file in LOCAL_RAW_DATA_DIR.rglob("*"): if file.is_file(): print(f" - {file.relative_to(LOCAL_RAW_DATA_DIR)}") - + except Exception as e: print(f"Error during download: {str(e)}") raise @@ -70,22 +70,37 @@ def run_preprocessing(): # Run torchrun command cmd = [ "torchrun", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING, + "--nnodes", + NUM_NODES, + "--nproc_per_node", + NUM_GPUS_PER_NODE_PREPROCESSING, PREPROCESSING_ENTRY_FILE_PATH, - "--model_path", MODEL_PATH, - "--data_merge_path", os.path.join(LOCAL_RAW_DATA_DIR, "merge_1_sample.txt"), - "--preprocess_video_batch_size", "1", - "--max_height", "480", - "--max_width", "832", - "--num_frames", "77", - "--dataloader_num_workers", "0", - "--output_dir", LOCAL_PREPROCESSED_DATA_DIR, - "--train_fps", "16", - "--samples_per_file", "1", - "--flush_frequency", "1", - "--video_length_tolerance_range", "5", - "--preprocess_task", "t2v", + "--model_path", + MODEL_PATH, + "--data_merge_path", + os.path.join(LOCAL_RAW_DATA_DIR, "merge_1_sample.txt"), + "--preprocess_video_batch_size", + "1", + "--max_height", + "480", + "--max_width", + "832", + "--num_frames", + "77", + "--dataloader_num_workers", + "0", + "--output_dir", + LOCAL_PREPROCESSED_DATA_DIR, + "--train_fps", + "16", + "--samples_per_file", + "1", + "--flush_frequency", + "1", + "--video_length_tolerance_range", + "5", + "--preprocess_task", + "t2v", ] process = subprocess.run(cmd, check=True) @@ -94,48 +109,87 @@ def run_preprocessing(): def run_training(): cmd = [ "torchrun", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING, + "--nnodes", + NUM_NODES, + "--nproc_per_node", + NUM_GPUS_PER_NODE_TRAINING, TRAINING_ENTRY_FILE_PATH, - "--model_path", MODEL_PATH, - "--inference_mode", "False", - "--pretrained_model_name_or_path", MODEL_PATH, - "--data_path", LOCAL_TRAINING_DATA_DIR, - "--validation_dataset_file", LOCAL_VALIDATION_DATASET_FILE, - "--train_batch_size", "1", - "--num_latent_t", "8", - "--num_gpus", NUM_GPUS_PER_NODE_TRAINING, - "--sp_size", NUM_GPUS_PER_NODE_TRAINING, - "--tp_size", "1", - "--hsdp_replicate_dim", "1", - "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING, - "--num_gpus", NUM_GPUS_PER_NODE_TRAINING, - "--train_sp_batch_size", "1", - "--dataloader_num_workers", "10", - "--gradient_accumulation_steps", "1", - "--max_train_steps", "901", - "--learning_rate", "5e-6", - "--mixed_precision", "bf16", - "--weight_only_checkpointing_steps", "6000", - "--training_state_checkpointing_steps", "6000", - "--validation_steps", "100", - "--validation_sampling_steps", "50", + "--model_path", + MODEL_PATH, + "--inference_mode", + "False", + "--pretrained_model_name_or_path", + MODEL_PATH, + "--data_path", + LOCAL_TRAINING_DATA_DIR, + "--validation_dataset_file", + LOCAL_VALIDATION_DATASET_FILE, + "--train_batch_size", + "1", + "--num_latent_t", + "8", + "--num_gpus", + NUM_GPUS_PER_NODE_TRAINING, + "--sp_size", + NUM_GPUS_PER_NODE_TRAINING, + "--tp_size", + "1", + "--hsdp_replicate_dim", + "1", + "--hsdp_shard_dim", + NUM_GPUS_PER_NODE_TRAINING, + "--num_gpus", + NUM_GPUS_PER_NODE_TRAINING, + "--train_sp_batch_size", + "1", + "--dataloader_num_workers", + "10", + "--gradient_accumulation_steps", + "1", + "--max_train_steps", + "901", + "--learning_rate", + "5e-6", + "--mixed_precision", + "bf16", + "--weight_only_checkpointing_steps", + "6000", + "--training_state_checkpointing_steps", + "6000", + "--validation_steps", + "100", + "--validation_sampling_steps", + "50", "--log_validation", - "--checkpoints_total_limit", "3", - "--ema_start_step", "0", - "--training_cfg_rate", "0.0", - "--output_dir", LOCAL_OUTPUT_DIR, - "--tracker_project_name", "wan_finetune_overfit_ci", - "--num_height", "480", - "--num_width", "832", - "--num_frames", "81", - "--validation_guidance_scale", "3.0", - "--num_euler_timesteps", "50", - "--multi_phased_distill_schedule", "4000-1", - "--weight_decay", "0.01", + "--checkpoints_total_limit", + "3", + "--ema_start_step", + "0", + "--training_cfg_rate", + "0.0", + "--output_dir", + LOCAL_OUTPUT_DIR, + "--tracker_project_name", + "wan_finetune_overfit_ci", + "--num_height", + "480", + "--num_width", + "832", + "--num_frames", + "81", + "--validation_guidance_scale", + "3.0", + "--num_euler_timesteps", + "50", + "--multi_phased_distill_schedule", + "4000-1", + "--weight_decay", + "0.01", "--not_apply_cfg_solver", - "--dit_precision", "fp32", - "--max_grad_norm", "1.0", + "--dit_precision", + "fp32", + "--max_grad_norm", + "1.0", ] print(f"Running training with command: {cmd}") @@ -154,7 +208,6 @@ def test_e2e_overfit_single_sample(): final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4") print(f"final_validation_video_file: {final_validation_video_file}") - # Ensure both files exist assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}" assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}" diff --git a/fastvideo/tests/ops/quantization/test_absmax_fp8.py b/fastvideo/tests/ops/quantization/test_absmax_fp8.py index 1f6f40e6f..462600b32 100644 --- a/fastvideo/tests/ops/quantization/test_absmax_fp8.py +++ b/fastvideo/tests/ops/quantization/test_absmax_fp8.py @@ -13,6 +13,7 @@ from fastvideo.models.utils import set_weight_attrs class TestAbsMaxFP8LinearMethod(unittest.TestCase): + def test_convert_scale_none(self): method = AbsMaxFP8LinearMethod() scale = method._convert_scale(None) @@ -81,17 +82,14 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase): output_size=2, params_dtype=torch.float16, ) - weight_fp16 = torch.tensor( - [[1.0, -2.0, 3.0], [4.0, 0.5, -1.5]], dtype=torch.float16 - ) + weight_fp16 = torch.tensor([[1.0, -2.0, 3.0], [4.0, 0.5, -1.5]], dtype=torch.float16) layer.weight.data = weight_fp16.to(dtype=torch.float8_e4m3fn) layer.scale_weight.data = torch.tensor([2.0, 3.0], dtype=torch.float32) layer.scale_input.data = torch.tensor([4.0], dtype=torch.float32) x = torch.tensor([[1.0, 2.0, -1.0]], dtype=torch.float16) expected = torch.nn.functional.linear( x * layer.scale_input.data.to(dtype=torch.float16), - weight_fp16 - * layer.scale_weight.data.to(dtype=torch.float16).unsqueeze(1), + weight_fp16 * layer.scale_weight.data.to(dtype=torch.float16).unsqueeze(1), ).to(dtype=torch.float16) output = method.apply(layer, x, bias=None) - assert_close(output, expected) \ No newline at end of file + assert_close(output, expected) diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_config.py b/fastvideo/tests/ops/quantization/test_nvfp4_config.py index 9f79980d5..a07fb7eb5 100644 --- a/fastvideo/tests/ops/quantization/test_nvfp4_config.py +++ b/fastvideo/tests/ops/quantization/test_nvfp4_config.py @@ -44,8 +44,7 @@ def test_nvfp4_kernel_call_raises_clear_error_without_flashinfer(monkeypatch): ImportError when flashinfer is missing, not a confusing AttributeError or NameError.""" # Stage a fake flashinfer that fails on import. - monkeypatch.setitem(sys.modules, "flashinfer", - _raise_module_on_import("flashinfer")) + monkeypatch.setitem(sys.modules, "flashinfer", _raise_module_on_import("flashinfer")) sys.modules.pop("fastvideo.layers.quantization.nvfp4_config", None) from fastvideo.layers.quantization.nvfp4_config import _require_flashinfer with pytest.raises(ImportError, match="flashinfer"): @@ -71,8 +70,7 @@ def test_coerce_fp4_input_dtype_casts_and_rejects(): cast to bf16, and non-floating inputs are rejected fast.""" import torch - from fastvideo.layers.quantization.nvfp4_config import ( - _coerce_fp4_input_dtype) + from fastvideo.layers.quantization.nvfp4_config import (_coerce_fp4_input_dtype) # bf16 / fp16 pass through untouched. bf16 = torch.zeros(4, 8, dtype=torch.bfloat16) @@ -81,10 +79,8 @@ def test_coerce_fp4_input_dtype_casts_and_rejects(): assert _coerce_fp4_input_dtype(fp16) is fp16 # Other floating dtypes (fp32 from an unfused norm, fp64) -> bf16. - assert _coerce_fp4_input_dtype( - torch.zeros(4, 8, dtype=torch.float32)).dtype is torch.bfloat16 - assert _coerce_fp4_input_dtype( - torch.zeros(4, 8, dtype=torch.float64)).dtype is torch.bfloat16 + assert _coerce_fp4_input_dtype(torch.zeros(4, 8, dtype=torch.float32)).dtype is torch.bfloat16 + assert _coerce_fp4_input_dtype(torch.zeros(4, 8, dtype=torch.float64)).dtype is torch.bfloat16 # Non-floating inputs are a real error, not silently cast. for bad in (torch.int32, torch.int64, torch.bool): diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_ltx2_wiring.py b/fastvideo/tests/ops/quantization/test_nvfp4_ltx2_wiring.py index d25ed912a..82dab3b78 100644 --- a/fastvideo/tests/ops/quantization/test_nvfp4_ltx2_wiring.py +++ b/fastvideo/tests/ops/quantization/test_nvfp4_ltx2_wiring.py @@ -51,11 +51,10 @@ def test_self_attention_to_q_k_v_out_get_nvfp4_method() -> None: attn = _matched_attn() for attr in ("to_q", "to_k", "to_v"): linear = getattr(attn, attr) - assert isinstance(linear, ReplicatedLinear), ( - f"{attr} must be ReplicatedLinear, got {type(linear).__name__}") - assert isinstance(linear.quant_method, NVFP4QuantizeMethod), ( - f"{attr}.quant_method must be NVFP4QuantizeMethod, got " - f"{type(linear.quant_method).__name__}") + assert isinstance(linear, ReplicatedLinear), (f"{attr} must be ReplicatedLinear, got {type(linear).__name__}") + assert isinstance(linear.quant_method, + NVFP4QuantizeMethod), (f"{attr}.quant_method must be NVFP4QuantizeMethod, got " + f"{type(linear.quant_method).__name__}") out_linear = attn.to_out[0] assert isinstance(out_linear, ReplicatedLinear) assert isinstance(out_linear.quant_method, NVFP4QuantizeMethod) @@ -153,18 +152,16 @@ def test_basic_av_block_propagates_quant_config_to_all_children() -> None: block.attn1.to_out[0]: "ltx2.blocks.3.attn1.to_out", block.attn2.to_q: "ltx2.blocks.3.attn2.to_q", block.attn2.to_out[0]: "ltx2.blocks.3.attn2.to_out", - block.audio_to_video_attn.to_q: - "ltx2.blocks.3.audio_to_video_attn.to_q", - block.video_to_audio_attn.to_v: - "ltx2.blocks.3.video_to_audio_attn.to_v", + block.audio_to_video_attn.to_q: "ltx2.blocks.3.audio_to_video_attn.to_q", + block.video_to_audio_attn.to_v: "ltx2.blocks.3.video_to_audio_attn.to_v", block.ff.net[0].proj: "ltx2.blocks.3.ffn.fc_in", block.ff.net[2]: "ltx2.blocks.3.ffn.fc_out", } for linear, expected_prefix in nvfp4_expectations.items(): assert isinstance(linear, ReplicatedLinear) - assert isinstance(linear.quant_method, NVFP4QuantizeMethod), ( - f"{expected_prefix} expected NVFP4QuantizeMethod, got " - f"{type(linear.quant_method).__name__}") + assert isinstance(linear.quant_method, + NVFP4QuantizeMethod), (f"{expected_prefix} expected NVFP4QuantizeMethod, got " + f"{type(linear.quant_method).__name__}") assert linear.quant_method.layer_prefix == expected_prefix # Confirm the non-quantized projections still get the @@ -209,14 +206,14 @@ def test_qat_train_uses_the_ltx2_deployment_target_profile() -> None: ) deployed_names = { - name for name, module in deployed.named_modules() - if isinstance(getattr(module, "quant_method", None), - NVFP4QuantizeMethod) + name + for name, module in deployed.named_modules() + if isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod) } qat_names = { - name for name, module in qat.named_modules() - if isinstance(getattr(module, "quant_method", None), - NVFP4QATTrainQuantizeMethod) + name + for name, module in qat.named_modules() + if isinstance(getattr(module, "quant_method", None), NVFP4QATTrainQuantizeMethod) } assert qat_names == deployed_names assert len(qat_names) == 12 diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_purge.py b/fastvideo/tests/ops/quantization/test_nvfp4_purge.py index 520910802..385cc4756 100644 --- a/fastvideo/tests/ops/quantization/test_nvfp4_purge.py +++ b/fastvideo/tests/ops/quantization/test_nvfp4_purge.py @@ -46,8 +46,7 @@ class _FakeLinear(nn.Module): def _fake_quantize(weight, global_sf, sfLayout=None, do_shuffle=False): out_dim, in_dim = weight.shape[0], weight.shape[-1] - return (torch.zeros(out_dim, in_dim // 2, dtype=torch.int8), torch.zeros(out_dim, in_dim // 16, - dtype=torch.uint8)) + return (torch.zeros(out_dim, in_dim // 2, dtype=torch.int8), torch.zeros(out_dim, in_dim // 16, dtype=torch.uint8)) @pytest.fixture(autouse=True) diff --git a/fastvideo/tests/ops/test_rotary_pos_embed_cache.py b/fastvideo/tests/ops/test_rotary_pos_embed_cache.py index 74f182ba9..2c4642471 100644 --- a/fastvideo/tests/ops/test_rotary_pos_embed_cache.py +++ b/fastvideo/tests/ops/test_rotary_pos_embed_cache.py @@ -76,8 +76,7 @@ def test_many_identical_calls_keep_single_entry(): ((4, 8, 8), 1536, 12, torch.float64, False), ], ) -def test_cached_matches_fresh_recompute(rope_sizes, hidden_size, heads_num, - dtype, use_real): +def test_cached_matches_fresh_recompute(rope_sizes, hidden_size, heads_num, dtype, use_real): """Cached tensors are bitwise-equal to a fresh uncached recompute.""" cos_cached, sin_cached = _call(rope_sizes=rope_sizes, hidden_size=hidden_size, @@ -97,12 +96,36 @@ def test_cached_matches_fresh_recompute(rope_sizes, hidden_size, heads_num, @pytest.mark.parametrize( "kwargs_a,kwargs_b", [ - ({"rope_sizes": (21, 30, 52)}, {"rope_sizes": (21, 45, 80)}), - ({"dtype": torch.float64}, {"dtype": torch.float32}), - ({"use_real": True}, {"use_real": False}), - ({"start_frame": 0}, {"start_frame": 3}), - ({"rope_theta": 10000.0}, {"rope_theta": 5000.0}), - ({"shard_dim": 0}, {"shard_dim": 1}), + ({ + "rope_sizes": (21, 30, 52) + }, { + "rope_sizes": (21, 45, 80) + }), + ({ + "dtype": torch.float64 + }, { + "dtype": torch.float32 + }), + ({ + "use_real": True + }, { + "use_real": False + }), + ({ + "start_frame": 0 + }, { + "start_frame": 3 + }), + ({ + "rope_theta": 10000.0 + }, { + "rope_theta": 5000.0 + }), + ({ + "shard_dim": 0 + }, { + "shard_dim": 1 + }), ], ) def test_distinct_args_create_distinct_entries(kwargs_a, kwargs_b): @@ -183,7 +206,7 @@ def test_cache_hit_refreshes_recency(): for frame in range(1, _ROTARY_POS_EMBED_CACHE_MAXSIZE): _call(rope_sizes=(2, 2, 2), start_frame=frame) assert len(_ROTARY_POS_EMBED_CACHE) == _ROTARY_POS_EMBED_CACHE_MAXSIZE - _call(rope_sizes=(2, 2, 2), start_frame=0) # hit -> frame 0 becomes most recent + _call(rope_sizes=(2, 2, 2), start_frame=0) # hit -> frame 0 becomes most recent _call(rope_sizes=(2, 2, 2), start_frame=99) # miss -> evicts now-oldest (frame 1) surviving = {key[-2] for key in _ROTARY_POS_EMBED_CACHE} assert 0 in surviving diff --git a/fastvideo/tests/performance/compare_baseline.py b/fastvideo/tests/performance/compare_baseline.py index 07a6f5ce8..a6fa9b37b 100644 --- a/fastvideo/tests/performance/compare_baseline.py +++ b/fastvideo/tests/performance/compare_baseline.py @@ -186,11 +186,7 @@ def _record_metadata(run_source: str, result: dict[str, Any]) -> dict[str, Any]: def _identity_metadata(result: dict[str, Any]) -> dict[str, Any]: - metadata = { - key: result[key] - for key in IDENTITY_KEYS - if key in result and result[key] is not None - } + metadata = {key: result[key] for key in IDENTITY_KEYS if key in result and result[key] is not None} recipe = result.get("recipe") benchmark = recipe.get("benchmark") if isinstance(recipe, dict) else None if isinstance(benchmark, dict): @@ -207,10 +203,7 @@ def _comparison_identity_filters(record: dict[str, Any]) -> dict[str, str]: ] if missing: raise ValueError("Performance record missing required comparison identity fields: " + ", ".join(missing)) - return { - key: str(record[key]) - for key in COMPARISON_IDENTITY_KEYS - } + return {key: str(record[key]) for key in COMPARISON_IDENTITY_KEYS} def _record_uses_v2_identity(record: dict[str, Any]) -> bool: @@ -220,10 +213,7 @@ def _record_uses_v2_identity(record: dict[str, Any]) -> bool: def _recipe_cohort_filters(record: dict[str, Any]) -> dict[str, str]: identity = _comparison_identity_filters(record) - return { - key: identity[key] - for key in ("workload_id", "variant_id", "benchmark_version") - } + return {key: identity[key] for key in ("workload_id", "variant_id", "benchmark_version")} def _comparison_identity_filters_or_none(record: dict[str, Any]) -> dict[str, str] | None: @@ -248,10 +238,8 @@ def _recipe_mismatch_records( current_recipe = str(record.get("recipe_fingerprint")) return [ item for item in cohort_records - if _record_can_authorize_recipe_mismatch(item) - and item.get("recipe_fingerprint") is not None - and str(item.get("recipe_fingerprint")).strip() - and str(item.get("recipe_fingerprint")) != current_recipe + if _record_can_authorize_recipe_mismatch(item) and item.get("recipe_fingerprint") is not None + and str(item.get("recipe_fingerprint")).strip() and str(item.get("recipe_fingerprint")) != current_recipe ] @@ -364,10 +352,8 @@ def _baseline_metric(records: list[dict[str, Any]], key: str) -> float | None: def _metric_policy_summary(policy: MetricPolicy) -> str: gated = "gated" if policy.gated else "info" - return ( - f"{gated}, >{policy.threshold_percent * 100:.1f}% " - f"and >{policy.threshold_absolute:.{policy.precision}f}" - ) + return (f"{gated}, >{policy.threshold_percent * 100:.1f}% " + f"and >{policy.threshold_absolute:.{policy.precision}f}") def _check_regressions( @@ -385,14 +371,12 @@ def _check_regressions( delta = regression_delta(policy, curr, baseline) if delta is None or not delta.regressed: continue - failures.append( - f"{current['model_id']} {policy.key} regressed by " - f"{delta.percent * 100:.1f}% and " - f"{delta.absolute:.{policy.precision}f} " - f"(current={curr:.{policy.precision}f}, " - f"baseline_median={baseline:.{policy.precision}f}, " - f"threshold={_metric_policy_summary(policy)})" - ) + failures.append(f"{current['model_id']} {policy.key} regressed by " + f"{delta.percent * 100:.1f}% and " + f"{delta.absolute:.{policy.precision}f} " + f"(current={curr:.{policy.precision}f}, " + f"baseline_median={baseline:.{policy.precision}f}, " + f"threshold={_metric_policy_summary(policy)})") return failures @@ -408,10 +392,8 @@ def _fixed_threshold_failures(raw_result: dict[str, Any]) -> list[str]: threshold = safe_float(thresholds.get(threshold_key)) if current is None or threshold is None or current <= threshold: continue - failures.append( - f"{raw_result.get('benchmark_id', 'unknown')} {result_key} exceeded fixed threshold " - f"(current={current:.3f}, threshold={threshold:.3f})" - ) + failures.append(f"{raw_result.get('benchmark_id', 'unknown')} {result_key} exceeded fixed threshold " + f"(current={current:.3f}, threshold={threshold:.3f})") return failures @@ -427,8 +409,7 @@ def _recipe_mismatch_failure( ) -> str: seen = sorted({ str(item.get("recipe_fingerprint")) - for item in recipe_mismatch_records - if item.get("recipe_fingerprint") is not None + for item in recipe_mismatch_records if item.get("recipe_fingerprint") is not None }) suffix = f"; existing recipe_fingerprint values: {', '.join(seen)}" if seen else "" return (f"{record['model_id']} recipe_fingerprint={record.get('recipe_fingerprint')} " @@ -465,10 +446,8 @@ def _evaluate_record_comparison( if not baseline_records: if identity_filters is not None: - return [], STATUS_CALIBRATION_NEEDED, ( - "No baseline found for exact comparable identity" - f"{_format_identity_filters(identity_filters)}" - ) + return [], STATUS_CALIBRATION_NEEDED, ("No baseline found for exact comparable identity" + f"{_format_identity_filters(identity_filters)}") return [], STATUS_PASS, f"No legacy baseline for {record['model_id']} on {record['gpu_type']}" @@ -505,11 +484,7 @@ def _build_summary_row( for policy in metric_policies: curr = safe_float(record.get(policy.key)) baseline = _baseline_metric(baseline_records, policy.key) - delta = ( - regression_delta(policy, curr, baseline) - if curr is not None and baseline is not None - else None - ) + delta = (regression_delta(policy, curr, baseline) if curr is not None and baseline is not None else None) regression = None if delta is None else delta.percent * 100.0 absolute_delta = None if delta is None else delta.absolute metric_values[policy.key] = { @@ -572,11 +547,8 @@ def _build_markdown_summary( f"{_compact_value(values['base'], policy.precision)}") worst_reg = ("n/a" if row["worst_regression_pct"] is None else f"{row['worst_regression_pct']:.1f}%") - exceeded_metrics = ( - ", ".join(row["threshold_exceeded_metrics"]) - if row["threshold_exceeded_metrics"] - else "none" - ) + exceeded_metrics = (", ".join(row["threshold_exceeded_metrics"]) + if row["threshold_exceeded_metrics"] else "none") failing_metrics = ", ".join(row["failing_metrics"]) if row["failing_metrics"] else "none" status = row["status"] status_detail = row["status_reason"] @@ -654,9 +626,8 @@ def main() -> int: recipe_mismatch_records: list[dict[str, Any]] = [] try: if identity_filters is None: - record["baseline_status"] = ( - "invalid_identity" if _record_uses_v2_identity(record) else "skipped_missing_identity" - ) + record["baseline_status"] = ("invalid_identity" + if _record_uses_v2_identity(record) else "skipped_missing_identity") else: baseline_records = load_records_for_identity( TRACKING_ROOT, @@ -674,9 +645,8 @@ def main() -> int: successful_only=True, ) recipe_mismatch_records = _recipe_mismatch_records(record, recipe_cohort_records) - record["baseline_status"] = ( - "recipe_mismatch" if recipe_mismatch_records else "initialized_new_cohort" - ) + record["baseline_status"] = ("recipe_mismatch" + if recipe_mismatch_records else "initialized_new_cohort") failures, comparison_status, comparison_status_reason = _evaluate_record_comparison( record, @@ -697,10 +667,8 @@ def main() -> int: record["comparison_status_reason"] = comparison_status_reason record["success"] = not failures - record["baseline_eligible"] = ( - identity_filters is not None - and _is_baseline_eligible(record["run_source"], record["success"], comparison_status) - ) + record["baseline_eligible"] = (identity_filters is not None and _is_baseline_eligible( + record["run_source"], record["success"], comparison_status)) all_failures.extend(failures) print(f"{record['model_id']} comparison status: " diff --git a/fastvideo/tests/performance/dashboard.py b/fastvideo/tests/performance/dashboard.py index 2364be364..f6f770f79 100644 --- a/fastvideo/tests/performance/dashboard.py +++ b/fastvideo/tests/performance/dashboard.py @@ -79,10 +79,7 @@ def _cohort_detail(record: object) -> str: def _group_record(record: object) -> dict[str, str]: - return { - key: _cohort_value(record.get(key)) - for key in ("model_id", "gpu_type", *COMPARISON_COHORT_KEYS) - } + return {key: _cohort_value(record.get(key)) for key in ("model_id", "gpu_type", *COMPARISON_COHORT_KEYS)} def _has_complete_v2_identity(record: object) -> bool: @@ -101,6 +98,7 @@ def _dashboard_frame(df: pd.DataFrame) -> pd.DataFrame: dashboard_df["_cohort_gpu_type"] = dashboard_df["gpu_type"].where(~uses_v2_identity, "") return dashboard_df + # ----------------------------- # 1. Grouping # ----------------------------- @@ -108,6 +106,7 @@ def group_data(df: pd.DataFrame): # Group by the same comparison cohort used by baseline gating. return _dashboard_frame(df).groupby(list(GROUP_KEYS), dropna=False) + # ----------------------------- # 2. Plot builder # ----------------------------- @@ -158,7 +157,10 @@ def build_plots(df: pd.DataFrame) -> tuple[list, list[dict[str, object]]]: markers=True, hover_data=["model_id", "gpu_type", "config_id", "commit_sha", *COMPARISON_COHORT_KEYS], title=f"{model_id} | {gpu_type} | {cohort_title} | {cohort_detail} | {metric}", - labels={"timestamp": "Time", metric: metric}, + labels={ + "timestamp": "Time", + metric: metric + }, ) figs.append(fig) @@ -177,25 +179,23 @@ def render_skipped_metrics(skipped_metrics: list[dict[str, object]]) -> str: "", ] for item in skipped_metrics: - rows.append( - "" - f"{escape(str(item['model_id']))}" - f"{escape(str(item['gpu_type']))}" - f"{escape(str(item['cohort']))}
{escape(str(item['cohort_detail']))}" - f"{escape(str(item['metric']))}" - f"{item['records']}" - f"{item['non_null']}" - f"{escape(str(item['reason']))}" - "" - ) + rows.append("" + f"{escape(str(item['model_id']))}" + f"{escape(str(item['gpu_type']))}" + f"{escape(str(item['cohort']))}
{escape(str(item['cohort_detail']))}" + f"{escape(str(item['metric']))}" + f"{item['records']}" + f"{item['non_null']}" + f"{escape(str(item['reason']))}" + "") rows.extend(["", ""]) return "\n".join(rows) + # ----------------------------- # 3. Render HTML dashboard # ----------------------------- -def render_html(figs: list, skipped_metrics: list[dict[str, object]], - days: int) -> str: +def render_html(figs: list, skipped_metrics: list[dict[str, object]], days: int) -> str: html_parts = [ "", "", @@ -220,6 +220,7 @@ def render_html(figs: list, skipped_metrics: list[dict[str, object]], html_parts.append("") return "\n".join(html_parts) + # ----------------------------- # 5. Main # ----------------------------- @@ -263,5 +264,6 @@ def main() -> None: print(f"Dashboard generated: {output_file}") + if __name__ == "__main__": main() diff --git a/fastvideo/tests/performance/identity.py b/fastvideo/tests/performance/identity.py index b222e4f29..4732620c3 100644 --- a/fastvideo/tests/performance/identity.py +++ b/fastvideo/tests/performance/identity.py @@ -37,14 +37,14 @@ _PROFILE_ENV_VARS = ( _PACKAGE_DISTRIBUTIONS = { "fastvideo_kernel": ("fastvideo-kernel", "fastvideo_kernel"), "flash_attn": ("flash-attn", "flash_attn"), - "flash_attn_4": ("flash-attn-4",), - "flash_attention_fp4": ("flash-attention-fp4",), + "flash_attn_4": ("flash-attn-4", ), + "flash_attention_fp4": ("flash-attention-fp4", ), "flashinfer": ("flashinfer-python", "flashinfer"), "nvidia_cutlass_dsl": ("nvidia-cutlass-dsl", "cutlass-dsl"), - "sageattention": ("sageattention",), - "sageattn3": ("sageattn3",), - "triton": ("triton",), - "xformers": ("xformers",), + "sageattention": ("sageattention", ), + "sageattn3": ("sageattn3", ), + "triton": ("triton", ), + "xformers": ("xformers", ), } @@ -72,10 +72,7 @@ def recipe_fingerprint(recipe: Mapping[str, Any]) -> str: # upstream repo commit, including model-card edits with unchanged # weights. The fingerprint is pinned to declared inputs only; the # resolved values stay in the stored recipe for auditability. - pruned = { - key: (dict(value) if isinstance(value, Mapping) else value) - for key, value in recipe.items() - } + pruned = {key: (dict(value) if isinstance(value, Mapping) else value) for key, value in recipe.items()} model = pruned.get("model") if isinstance(model, dict): model.pop("resolved_revision", None) @@ -114,16 +111,10 @@ def environment_fingerprint(metadata: Mapping[str, Any]) -> str: def benchmark_identity_from_config(cfg: Mapping[str, Any]) -> dict[str, Any]: - missing = [ - key for key in REQUIRED_BENCHMARK_IDENTITY_KEYS - if _none_if_empty(cfg.get(key)) is None - ] + missing = [key for key in REQUIRED_BENCHMARK_IDENTITY_KEYS if _none_if_empty(cfg.get(key)) is None] if missing: raise ValueError("Benchmark config missing required identity fields: " + ", ".join(missing)) - return { - key: cfg[key] - for key in REQUIRED_BENCHMARK_IDENTITY_KEYS - } + return {key: cfg[key] for key in REQUIRED_BENCHMARK_IDENTITY_KEYS} def build_recipe_from_benchmark_config( @@ -276,8 +267,7 @@ def software_profile( "cuda": _major_minor(cuda_version or torch.version.cuda), "packages": { name: str(version) - for name, version in sorted(versions.items()) - if version + for name, version in sorted(versions.items()) if version }, } @@ -310,17 +300,16 @@ def environment_metadata( }, "env": { key: source_env.get(key) - for key in _PROFILE_ENV_VARS - if source_env.get(key) is not None + for key in _PROFILE_ENV_VARS if source_env.get(key) is not None }, "packages": { key: value - for key, value in sorted(full_package_versions.items()) - if value + for key, value in sorted(full_package_versions.items()) if value }, - "hardware_profile": dict(hardware) if hardware is not None else hardware_profile(), - "software_profile": dict(software) if software is not None else software_profile( - package_versions=full_package_versions), + "hardware_profile": + dict(hardware) if hardware is not None else hardware_profile(), + "software_profile": + dict(software) if software is not None else software_profile(package_versions=full_package_versions), } @@ -328,10 +317,7 @@ def _canonicalize(value: Any) -> Any: if dataclasses.is_dataclass(value) and not isinstance(value, type): return _canonicalize(dataclasses.asdict(value)) if isinstance(value, Mapping): - return { - str(key): _canonicalize(value[key]) - for key in sorted(value, key=lambda item: str(item)) - } + return {str(key): _canonicalize(value[key]) for key in sorted(value, key=lambda item: str(item))} if isinstance(value, (set, frozenset)): return [_canonicalize(item) for item in sorted(value, key=lambda item: str(item))] if isinstance(value, tuple): diff --git a/fastvideo/tests/performance/seed_baseline.py b/fastvideo/tests/performance/seed_baseline.py index 3f9cc48a4..af45fe520 100644 --- a/fastvideo/tests/performance/seed_baseline.py +++ b/fastvideo/tests/performance/seed_baseline.py @@ -271,16 +271,14 @@ def _validate_no_existing_seed( ) -> None: identity = _source_identity(record) if load_records_for_identity( - records_root, - identity, - last_n=1, - successful_only=True, - baseline_eligible_only=True, + records_root, + identity, + last_n=1, + successful_only=True, + baseline_eligible_only=True, ): - raise ValueError( - f"{source} already has a baseline-eligible record for the exact comparable identity; " - "the CALIBRATION_NEEDED source artifact is stale" - ) + raise ValueError(f"{source} already has a baseline-eligible record for the exact comparable identity; " + "the CALIBRATION_NEEDED source artifact is stale") cohort_records = load_records_for_identity( records_root, @@ -290,10 +288,8 @@ def _validate_no_existing_seed( mismatches = _recipe_mismatch_records(record, cohort_records) if mismatches: recipes = sorted({str(item["recipe_fingerprint"]) for item in mismatches}) - raise ValueError( - f"{source} already has trusted records for another recipe in this workload, " - f"variant, and benchmark version: {', '.join(recipes)}" - ) + raise ValueError(f"{source} already has trusted records for another recipe in this workload, " + f"variant, and benchmark version: {', '.join(recipes)}") def _repository_descriptor() -> dict[str, str]: @@ -308,7 +304,9 @@ def _reservation_path(staging_root: str, identity: dict[str, str]) -> str: payload = json.dumps({ "identity": identity, "repository": _repository_descriptor(), - }, sort_keys=True, separators=(",", ":")) + }, + sort_keys=True, + separators=(",", ":")) digest = hashlib.sha256(payload.encode("utf-8")).hexdigest() return os.path.join(staging_root, ".seed-reservations", digest) @@ -319,10 +317,8 @@ def _reserve_staging_identity(staging_root: str, identity: dict[str, str]) -> st try: os.mkdir(reservation, 0o700) except FileExistsError as exc: - raise ValueError( - "staging already has a reservation for the exact comparable identity; " - "reuse or explicitly clean the existing preparation" - ) from exc + raise ValueError("staging already has a reservation for the exact comparable identity; " + "reuse or explicitly clean the existing preparation") from exc return reservation @@ -375,15 +371,20 @@ def _write_preparation_manifest( "sha256": hashlib.sha256(blob).hexdigest(), }) manifest = { - "manifest_schema_version": 1, - "repository": _repository_descriptor(), - "identity": identity, + "manifest_schema_version": + 1, + "repository": + _repository_descriptor(), + "identity": + identity, "recipe_cohort": { key: identity[key] for key in ("workload_id", "variant_id", "benchmark_version") }, - "staging_root": os.path.realpath(os.path.abspath(staging_root)), - "source_records": source_manifest_entries, + "staging_root": + os.path.realpath(os.path.abspath(staging_root)), + "source_records": + source_manifest_entries, "prepared_records": [{ "path": os.path.realpath(os.path.abspath(path)), "sha256": _file_sha256(path), @@ -408,13 +409,11 @@ def _validate_prepared_seed(record: dict[str, Any]) -> None: raise ValueError("prepared records must retain CALIBRATION_NEEDED source provenance") if not str(record.get("baseline_seed_reason") or "").strip(): raise ValueError("prepared records must retain a non-empty approval rationale") - if ( - record.get("baseline_seed_source_success") is not True - or record.get("baseline_seed_source_run_source") != "scheduled_main" - or record.get("baseline_seed_source_branch") != "main" - or record.get("baseline_seed_source_test_scope") != "full" - or _truthy_pr_number(record.get("baseline_seed_source_pr_number")) - ): + if (record.get("baseline_seed_source_success") is not True + or record.get("baseline_seed_source_run_source") != "scheduled_main" + or record.get("baseline_seed_source_branch") != "main" + or record.get("baseline_seed_source_test_scope") != "full" + or _truthy_pr_number(record.get("baseline_seed_source_pr_number"))): raise ValueError("prepared records must retain trusted scheduled-main source provenance") @@ -515,8 +514,8 @@ def upload_prepared_seed_manifest(manifest_path: str) -> str: ordered_sources = _order_sources_by_timestamp(source_paths, source_records) records_by_index = sorted(records, key=lambda record: record["baseline_seed_batch_index"]) for index, ((source_path, source_record), record) in enumerate( - zip(ordered_sources, records_by_index, strict=True), - start=1, + zip(ordered_sources, records_by_index, strict=True), + start=1, ): expected = build_baseline_seed_record( source_record, @@ -589,10 +588,8 @@ def _order_sources_by_timestamp( timestamp = None if timestamp is None: - print( - f"Warning: source artifact {index} has a missing or unparsable timestamp; " - "preserving its input order after timestamped sources." - ) + print(f"Warning: source artifact {index} has a missing or unparsable timestamp; " + "preserving its input order after timestamped sources.") undated.append((source_path, record)) continue @@ -621,11 +618,8 @@ def _validate_batch_consistency( print(f"Source batch consistency (maximum regression {max_regression * 100:.1f}%):") failures = [] for policy in DEFAULT_METRIC_POLICIES: - values = [ - (index, value) - for index, record in enumerate(records, start=1) - if (value := safe_float(record.get(policy.key))) is not None - ] + values = [(index, value) for index, record in enumerate(records, start=1) + if (value := safe_float(record.get(policy.key))) is not None] if len(values) < 2: continue @@ -640,16 +634,14 @@ def _validate_batch_consistency( else: regression = (batch_median - value) / batch_median regressions.append((source_index, regression)) - print( - f" {policy.key} source={source_index} value={value:.6f} " - f"median={batch_median:.6f} regression={regression * 100:.1f}%") + print(f" {policy.key} source={source_index} value={value:.6f} " + f"median={batch_median:.6f} regression={regression * 100:.1f}%") worst_source, worst_regression = max(regressions, key=lambda item: item[1]) print(f" {policy.key} worst regression: {worst_regression * 100:.1f}% (source={worst_source})") if worst_regression > max_regression: - failures.append( - f"source artifact {worst_source} {policy.key} regresses by " - f"{worst_regression * 100:.1f}% against the batch median") + failures.append(f"source artifact {worst_source} {policy.key} regresses by " + f"{worst_regression * 100:.1f}% against the batch median") if failures: raise ValueError("source batch exceeds maximum intra-batch regression: " + "; ".join(failures)) @@ -657,8 +649,7 @@ def _validate_batch_consistency( def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser( - description="Create reviewed baseline seed records from v2 CALIBRATION_NEEDED artifacts.", - ) + description="Create reviewed baseline seed records from v2 CALIBRATION_NEEDED artifacts.", ) parser.add_argument( "--source-result", dest="source_results", @@ -674,10 +665,8 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument( "--tracking-root", default=TRACKING_ROOT, - help=( - "Operator tracking-mirror path used only to enforce staging separation; " - "remote validation uses a fresh temporary snapshot." - ), + help=("Operator tracking-mirror path used only to enforce staging separation; " + "remote validation uses a fresh temporary snapshot."), ) parser.add_argument( "--staging-root", @@ -730,10 +719,8 @@ def main(argv: list[str] | None = None) -> int: seed_records.append((seed_record, suffix)) reservation = _reserve_staging_identity(args.staging_root, identity) - prepared_seeds = [ - (_seed_record_path(reservation, seed_record, suffix=suffix), seed_record, suffix) - for seed_record, suffix in seed_records - ] + prepared_seeds = [(_seed_record_path(reservation, seed_record, suffix=suffix), seed_record, suffix) + for seed_record, suffix in seed_records] prepared_paths = [seed_path for seed_path, _seed_record, _suffix in prepared_seeds] try: for seed_path, seed_record, suffix in prepared_seeds: diff --git a/fastvideo/tests/performance/test_benchmark_config_validation.py b/fastvideo/tests/performance/test_benchmark_config_validation.py index 61ee1e781..0979ffa3d 100644 --- a/fastvideo/tests/performance/test_benchmark_config_validation.py +++ b/fastvideo/tests/performance/test_benchmark_config_validation.py @@ -44,7 +44,9 @@ def test_v2_benchmark_config_identity_validates_and_is_preserved(): "workload_id": "wan-t2v", "variant_id": "1.3b-sp2", "benchmark_version": 2, - "quality_metadata": {"some": "data"}, + "quality_metadata": { + "some": "data" + }, } @@ -113,7 +115,11 @@ def test_optional_v2_metadata_fields_must_be_objects(): def test_regression_thresholds_validate_without_forcing_v2_schema(): cfg = { "benchmark_id": "legacy-benchmark", - "regression_thresholds": {"latency": {"threshold_percent": 0.1}}, + "regression_thresholds": { + "latency": { + "threshold_percent": 0.1 + } + }, } _validate_benchmark_config(cfg, "legacy.json") diff --git a/fastvideo/tests/performance/test_compare_baseline_policy.py b/fastvideo/tests/performance/test_compare_baseline_policy.py index 1af9a5a95..5de5b7316 100644 --- a/fastvideo/tests/performance/test_compare_baseline_policy.py +++ b/fastvideo/tests/performance/test_compare_baseline_policy.py @@ -114,7 +114,8 @@ def test_sync_marker_reuse_requires_the_exact_revision(tmp_path): "endpoint": hf_store.ENDPOINT, "repo_id": hf_store.HF_REPO_ID, "revision": "pinned-old-revision", - }), encoding="utf-8") + }), + encoding="utf-8") assert hf_store._sync_marker_matches_request(str(marker), "pinned-old-revision") is True assert hf_store._sync_marker_matches_request(str(marker), None) is False @@ -123,7 +124,8 @@ def test_sync_marker_reuse_requires_the_exact_revision(tmp_path): "endpoint": "https://different.example", "repo_id": hf_store.HF_REPO_ID, "revision": "pinned-old-revision", - }), encoding="utf-8") + }), + encoding="utf-8") assert hf_store._sync_marker_matches_request(str(marker), "pinned-old-revision") is False @@ -217,18 +219,19 @@ def test_invalid_regression_threshold_container_uses_defaults(): def test_boolean_regression_threshold_values_are_ignored(): - policies = resolve_metric_policies({ - "latency": { + policies = resolve_metric_policies( + {"latency": { "threshold_percent": True, "threshold_absolute": False, "gated": "false", - } - }) + }}) latency = next(policy for policy in policies if policy.key == "latency") assert latency.threshold_percent == 0.08 assert latency.threshold_absolute == 0.5 assert latency.gated is False + + def test_normalized_record_preserves_identity_metadata(monkeypatch): monkeypatch.setenv("PERF_RUN_SOURCE", "pr") raw = _raw_result() @@ -443,10 +446,8 @@ def test_comparison_identity_filters_keep_zero_version(): def test_baseline_eligibility_only_for_successful_scheduled_main(): assert compare_baseline._is_baseline_eligible("scheduled_main", True) is True assert compare_baseline._is_baseline_eligible("scheduled_main", False) is False - assert compare_baseline._is_baseline_eligible( - "scheduled_main", True, compare_baseline.STATUS_PASS) is True - assert compare_baseline._is_baseline_eligible( - "scheduled_main", False, compare_baseline.STATUS_PASS) is False + assert compare_baseline._is_baseline_eligible("scheduled_main", True, compare_baseline.STATUS_PASS) is True + assert compare_baseline._is_baseline_eligible("scheduled_main", False, compare_baseline.STATUS_PASS) is False assert compare_baseline._is_baseline_eligible("pr", True) is False assert compare_baseline._is_baseline_eligible("local", True) is False @@ -455,24 +456,18 @@ def test_latency_regression_requires_percent_and_absolute_floors(): baseline = [{"latency": 10.0}] current = {"model_id": "wan", "latency": 10.6} - percent_only = resolve_metric_policies({ - "latency": { - "threshold_percent": 0.05, - "threshold_absolute": 0.75, - } - }) - absolute_only = resolve_metric_policies({ - "latency": { - "threshold_percent": 0.10, - "threshold_absolute": 0.5, - } - }) - both = resolve_metric_policies({ - "latency": { - "threshold_percent": 0.05, - "threshold_absolute": 0.5, - } - }) + percent_only = resolve_metric_policies({"latency": { + "threshold_percent": 0.05, + "threshold_absolute": 0.75, + }}) + absolute_only = resolve_metric_policies({"latency": { + "threshold_percent": 0.10, + "threshold_absolute": 0.5, + }}) + both = resolve_metric_policies({"latency": { + "threshold_percent": 0.05, + "threshold_absolute": 0.5, + }}) assert compare_baseline._check_regressions(current, baseline, percent_only) == [] assert compare_baseline._check_regressions(current, baseline, absolute_only) == [] @@ -485,12 +480,10 @@ def test_latency_regression_requires_percent_and_absolute_floors(): def test_throughput_regression_uses_higher_is_better_direction(): baseline = [{"throughput": 10.0}] current = {"model_id": "wan", "throughput": 9.0} - policies = resolve_metric_policies({ - "throughput": { - "threshold_percent": 0.05, - "threshold_absolute": 0.5, - } - }) + policies = resolve_metric_policies({"throughput": { + "threshold_percent": 0.05, + "threshold_absolute": 0.5, + }}) failures = compare_baseline._check_regressions(current, baseline, policies) @@ -501,12 +494,10 @@ def test_throughput_regression_uses_higher_is_better_direction(): def test_memory_regression_uses_metric_specific_absolute_floor(): baseline = [{"memory": 10000.0}] current = {"model_id": "wan", "memory": 10600.0} - policies = resolve_metric_policies({ - "memory": { - "threshold_percent": 0.05, - "threshold_absolute": 256.0, - } - }) + policies = resolve_metric_policies({"memory": { + "threshold_percent": 0.05, + "threshold_absolute": 256.0, + }}) failures = compare_baseline._check_regressions(current, baseline, policies) @@ -517,12 +508,10 @@ def test_memory_regression_uses_metric_specific_absolute_floor(): def test_component_metric_can_gate_independently(): baseline = [{"dit_time_s": 8.0}] current = {"model_id": "wan", "dit_time_s": 8.6} - policies = resolve_metric_policies({ - "dit_time_s": { - "threshold_percent": 0.05, - "threshold_absolute": 0.25, - } - }) + policies = resolve_metric_policies({"dit_time_s": { + "threshold_percent": 0.05, + "threshold_absolute": 0.25, + }}) failures = compare_baseline._check_regressions(current, baseline, policies) @@ -533,13 +522,12 @@ def test_component_metric_can_gate_independently(): def test_informational_metric_remains_visible_without_failing(): baseline = [{"throughput": 10.0}] current = {"model_id": "wan", "gpu_type": "NVIDIA L40S", "throughput": 8.0} - policies = resolve_metric_policies({ - "throughput": { + policies = resolve_metric_policies( + {"throughput": { "threshold_percent": 0.01, "threshold_absolute": 0.01, "gated": False, - } - }) + }}) row = compare_baseline._build_summary_row(current, baseline, policies, False) @@ -757,8 +745,12 @@ def test_baseline_seed_rejects_untrusted_calibration_sources(monkeypatch, overri @pytest.mark.parametrize( ("overrides", "match"), [ - ({"result_schema_version": 1}, "normalized v2"), - ({"baseline_eligible": True}, "baseline_eligible=false"), + ({ + "result_schema_version": 1 + }, "normalized v2"), + ({ + "baseline_eligible": True + }, "baseline_eligible=false"), ], ) def test_baseline_seed_rejects_invalid_calibration_schema(overrides, match): @@ -866,14 +858,10 @@ def test_baseline_seed_orders_sources_by_original_timestamp(monkeypatch, tmp_pat manifest_path = next(staging_root.glob(".seed-reservations/*/manifest")) manifest = json.loads(manifest_path.read_text(encoding="utf-8")) - prepared = [ - json.loads(Path(entry["path"]).read_text(encoding="utf-8")) - for entry in manifest["prepared_records"] - ] + prepared = [json.loads(Path(entry["path"]).read_text(encoding="utf-8")) for entry in manifest["prepared_records"]] last_five = prepared[-5:] - assert [record["baseline_seed_source_timestamp"] for record in last_five] == [ - f"2026-06-{day:02d}T00:00:00+00:00" for day in range(2, 7) - ] + assert [record["baseline_seed_source_timestamp"] + for record in last_five] == [f"2026-06-{day:02d}T00:00:00+00:00" for day in range(2, 7)] assert not tracking_root.exists() assert len(list(staging_root.glob(".seed-reservations/*/manifest"))) == 1 @@ -898,6 +886,7 @@ def test_baseline_seed_rejects_existing_eligible_identity(monkeypatch, tmp_path) }) tracking_root = tmp_path / "tracking" staging_root = tmp_path / "staging" + def sync(local_dir, *, strict=False, revision=None): assert strict is True assert revision is None @@ -944,6 +933,7 @@ def test_baseline_seed_rejects_trusted_different_recipe(monkeypatch, tmp_path): }) tracking_root = tmp_path / "tracking" staging_root = tmp_path / "staging" + def sync(local_dir, *, strict=False, revision=None): assert strict is True assert revision is None @@ -1062,6 +1052,7 @@ def test_baseline_seed_upload_is_one_conditional_commit(monkeypatch, tmp_path): calls = {} class FakeApi: + def __init__(self, token): assert token == "hf-test" @@ -1130,6 +1121,7 @@ def test_baseline_seed_upload_rechecks_remote_state(monkeypatch, tmp_path): return local_dir class FakeApi: + def __init__(self, token): assert token == "hf-test" @@ -1151,6 +1143,7 @@ def test_baseline_seed_upload_preserves_manifest_on_cas_failure(monkeypatch, tmp manifest_path = _prepare_single_seed(tmp_path) class FakeApi: + def __init__(self, token): assert token == "hf-test" @@ -1184,10 +1177,16 @@ def test_baseline_seed_upload_rejects_staged_mutation(tmp_path): def test_baseline_seed_orders_invalid_timestamps_last_and_warns(capsys): source_paths = ["dated-later", "missing", "malformed", "dated-earlier"] records = [ - {"timestamp": "2026-06-02T00:00:00+00:00"}, + { + "timestamp": "2026-06-02T00:00:00+00:00" + }, {}, - {"timestamp": "not-a-timestamp"}, - {"timestamp": "2026-06-01T00:00:00+00:00"}, + { + "timestamp": "not-a-timestamp" + }, + { + "timestamp": "2026-06-01T00:00:00+00:00" + }, ] ordered = seed_baseline._order_sources_by_timestamp(source_paths, records) @@ -1340,7 +1339,7 @@ def test_baseline_seed_rejects_invalid_measurements_before_persistence( else: invalid_source[metric] = invalid_value - sources = (valid_source, invalid_source) if valid_prefix else (invalid_source,) + sources = (valid_source, invalid_source) if valid_prefix else (invalid_source, ) source_paths = [] for index, source in enumerate(sources, start=1): source_path = tmp_path / f"source_{index}.json" @@ -1740,9 +1739,8 @@ def test_scheduled_main_static_regression_does_not_contaminate_passing_record( normalized[record["model_id"]] = record assert normalized["regressed-benchmark"]["comparison_status"] == compare_baseline.STATUS_REGRESSION - assert "avg_generation_time_s exceeded fixed threshold" in normalized[ - "regressed-benchmark" - ]["comparison_status_reason"] + assert "avg_generation_time_s exceeded fixed threshold" in normalized["regressed-benchmark"][ + "comparison_status_reason"] assert normalized["regressed-benchmark"]["success"] is False assert normalized["regressed-benchmark"]["baseline_eligible"] is False assert normalized["passing-benchmark"]["comparison_status"] == compare_baseline.STATUS_PASS @@ -1756,7 +1754,9 @@ def test_unattributed_performance_pytest_failure_reports_infra_error(monkeypatch failures, status, reason = compare_baseline._evaluate_record_comparison( record, - [{"latency": 10.0}], + [{ + "latency": 10.0 + }], [], resolve_metric_policies(None), [], @@ -1923,8 +1923,8 @@ def test_legacy_record_lookup_still_uses_model_and_gpu(tmp_path): def test_non_pass_status_is_not_baseline_eligible(): - assert compare_baseline._is_baseline_eligible( - "scheduled_main", True, compare_baseline.STATUS_CALIBRATION_NEEDED) is False + assert compare_baseline._is_baseline_eligible("scheduled_main", True, + compare_baseline.STATUS_CALIBRATION_NEEDED) is False assert compare_baseline._is_baseline_eligible("pr", True, compare_baseline.STATUS_PASS) is False assert compare_baseline._is_baseline_eligible("local", True, compare_baseline.STATUS_PASS) is False diff --git a/fastvideo/tests/performance/test_dashboard_api.py b/fastvideo/tests/performance/test_dashboard_api.py index 649b1282d..2bbc12679 100644 --- a/fastvideo/tests/performance/test_dashboard_api.py +++ b/fastvideo/tests/performance/test_dashboard_api.py @@ -6,6 +6,7 @@ from fastvideo.performance_dashboard.api import PerformanceDataStore, create_app class FakeStore(PerformanceDataStore): + def __init__(self, records): super().__init__(tracking_root="/tmp/fake-fastvideo-perf-dashboard") self._records = records @@ -55,36 +56,37 @@ def _record(model_id, gpu_type, ts, commit, latency, throughput, success=True, * def test_summary_endpoint_returns_latest_group_status(): - app = create_app(FakeStore([ - _record( - "wan", - "NVIDIA L40S", - "2026-01-01T00:00:00+00:00", - "a" * 40, - 10.0, - 10.0, - workload_id="wan-t2v", - variant_id="1.3b-sp2", - benchmark_version=2, - recipe_fingerprint="recipe-a", - hardware_profile_id="hw-l40s", - software_profile_id="sw-cu130", - ), - _record( - "wan", - "NVIDIA L40S", - "2026-01-02T00:00:00+00:00", - "b" * 40, - 11.0, - 9.0, - workload_id="wan-t2v", - variant_id="1.3b-sp2", - benchmark_version=2, - recipe_fingerprint="recipe-a", - hardware_profile_id="hw-l40s", - software_profile_id="sw-cu130", - ), - ])) + app = create_app( + FakeStore([ + _record( + "wan", + "NVIDIA L40S", + "2026-01-01T00:00:00+00:00", + "a" * 40, + 10.0, + 10.0, + workload_id="wan-t2v", + variant_id="1.3b-sp2", + benchmark_version=2, + recipe_fingerprint="recipe-a", + hardware_profile_id="hw-l40s", + software_profile_id="sw-cu130", + ), + _record( + "wan", + "NVIDIA L40S", + "2026-01-02T00:00:00+00:00", + "b" * 40, + 11.0, + 9.0, + workload_id="wan-t2v", + variant_id="1.3b-sp2", + benchmark_version=2, + recipe_fingerprint="recipe-a", + hardware_profile_id="hw-l40s", + software_profile_id="sw-cu130", + ), + ])) client = TestClient(app) response = client.get("/api/performance/summary") @@ -104,10 +106,11 @@ def test_summary_endpoint_returns_latest_group_status(): def test_summary_status_is_independent_of_days_window(): - app = create_app(FakeStore([ - _record("wan", "NVIDIA L40S", "2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0), - _record("wan", "NVIDIA L40S", "2026-02-15T00:00:00+00:00", "b" * 40, 11.0, 9.0), - ])) + app = create_app( + FakeStore([ + _record("wan", "NVIDIA L40S", "2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0), + _record("wan", "NVIDIA L40S", "2026-02-15T00:00:00+00:00", "b" * 40, 11.0, 9.0), + ])) client = TestClient(app) narrow = client.get("/api/performance/summary", params={"days": 1}).json() @@ -122,30 +125,31 @@ def test_summary_status_is_independent_of_days_window(): def test_dashboard_endpoints_filter_and_return_run_source_metadata(): - app = create_app(FakeStore([ - _record( - "wan", - "NVIDIA L40S", - "2026-01-01T00:00:00+00:00", - "a" * 40, - 10.0, - 10.0, - run_source="pr", - pr_number="123", - branch="feature/perf", - baseline_eligible=False, - ), - _record( - "wan", - "NVIDIA L40S", - "2026-01-02T00:00:00+00:00", - "b" * 40, - 11.0, - 9.0, - run_source="scheduled_main", - baseline_eligible=True, - ), - ])) + app = create_app( + FakeStore([ + _record( + "wan", + "NVIDIA L40S", + "2026-01-01T00:00:00+00:00", + "a" * 40, + 10.0, + 10.0, + run_source="pr", + pr_number="123", + branch="feature/perf", + baseline_eligible=False, + ), + _record( + "wan", + "NVIDIA L40S", + "2026-01-02T00:00:00+00:00", + "b" * 40, + 11.0, + 9.0, + run_source="scheduled_main", + baseline_eligible=True, + ), + ])) client = TestClient(app) summary = client.get("/api/performance/summary", params={"run_source": "pr"}).json() @@ -162,10 +166,11 @@ def test_dashboard_endpoints_filter_and_return_run_source_metadata(): def test_records_and_trends_endpoints_filter_by_model_and_gpu(): - app = create_app(FakeStore([ - _record("wan", "NVIDIA L40S", "2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0), - _record("ltx", "NVIDIA A100", "2026-01-01T00:00:00+00:00", "b" * 40, 20.0, 5.0), - ])) + app = create_app( + FakeStore([ + _record("wan", "NVIDIA L40S", "2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0), + _record("ltx", "NVIDIA A100", "2026-01-01T00:00:00+00:00", "b" * 40, 20.0, 5.0), + ])) client = TestClient(app) records = client.get("/api/performance/records", params={"model_id": "wan"}).json() @@ -187,28 +192,29 @@ def test_v2_display_filters_keep_renamed_cohort_history(): "hardware_profile_id": "hw-l40s", "software_profile_id": "sw-cu130", } - app = create_app(FakeStore([ - _record( - "old-display", - "NVIDIA L40S old label", - "2026-01-01T00:00:00+00:00", - "a" * 40, - 10.0, - 10.0, - run_source="scheduled_main", - baseline_eligible=True, - **identity, - ), - _record( - "new-display", - "NVIDIA L40S", - "2026-01-02T00:00:00+00:00", - "b" * 40, - 11.0, - 9.0, - **identity, - ), - ])) + app = create_app( + FakeStore([ + _record( + "old-display", + "NVIDIA L40S old label", + "2026-01-01T00:00:00+00:00", + "a" * 40, + 10.0, + 10.0, + run_source="scheduled_main", + baseline_eligible=True, + **identity, + ), + _record( + "new-display", + "NVIDIA L40S", + "2026-01-02T00:00:00+00:00", + "b" * 40, + 11.0, + 9.0, + **identity, + ), + ])) client = TestClient(app) filters = {"model_id": "new-display", "gpu_type": "NVIDIA L40S"} diff --git a/fastvideo/tests/performance/test_dashboard_plotly.py b/fastvideo/tests/performance/test_dashboard_plotly.py index 8f9f4a338..7be8f2144 100644 --- a/fastvideo/tests/performance/test_dashboard_plotly.py +++ b/fastvideo/tests/performance/test_dashboard_plotly.py @@ -60,17 +60,14 @@ def test_group_data_ignores_v2_display_name_changes(): def test_group_data_fills_legacy_cohort_columns(): - df = pd.DataFrame([ - { - "model_id": model_id, - "gpu_type": "NVIDIA L40S", - "timestamp": "2026-01-01T00:00:00+00:00", - "commit_sha": "a" * 40, - "config_id": "aaaaaaa", - "latency": 10.0, - } - for model_id in ("legacy-wan", "legacy-ltx") - ]) + df = pd.DataFrame([{ + "model_id": model_id, + "gpu_type": "NVIDIA L40S", + "timestamp": "2026-01-01T00:00:00+00:00", + "commit_sha": "a" * 40, + "config_id": "aaaaaaa", + "latency": 10.0, + } for model_id in ("legacy-wan", "legacy-ltx")]) groups = list(dashboard.group_data(df)) @@ -97,10 +94,11 @@ def test_group_data_keeps_partial_v2_identity_scoped_by_display_metadata(): groups = list(dashboard.group_data(df)) assert len(groups) == 2 - assert {(key[1], key[2]) for key, _group in groups} == { - ("wan", "NVIDIA L40S"), - ("ltx", "NVIDIA H100"), - } + assert {(key[1], key[2]) + for key, _group in groups} == { + ("wan", "NVIDIA L40S"), + ("ltx", "NVIDIA H100"), + } def test_build_plots_labels_distinct_cohorts(): diff --git a/fastvideo/tests/performance/test_dashboard_service.py b/fastvideo/tests/performance/test_dashboard_service.py index 806cccd22..cb76b2ad0 100644 --- a/fastvideo/tests/performance/test_dashboard_service.py +++ b/fastvideo/tests/performance/test_dashboard_service.py @@ -98,13 +98,11 @@ def test_build_latest_summary_requires_absolute_floor_for_computed_regression(): "b" * 40, 10.6, 10.0, - regression_thresholds={ - "latency": { - "threshold_percent": 0.05, - "threshold_absolute": 0.75, - "gated": True, - } - }, + regression_thresholds={"latency": { + "threshold_percent": 0.05, + "threshold_absolute": 0.75, + "gated": True, + }}, ), ] @@ -125,13 +123,11 @@ def test_build_latest_summary_separates_informational_threshold_crossing(): "b" * 40, 10.6, 10.0, - regression_thresholds={ - "latency": { - "threshold_percent": 0.05, - "threshold_absolute": 0.5, - "gated": False, - } - }, + regression_thresholds={"latency": { + "threshold_percent": 0.05, + "threshold_absolute": 0.5, + "gated": False, + }}, ), ] @@ -182,7 +178,6 @@ def test_build_latest_summary_run_source_filter_excludes_future_baselines(): assert rows[0]["metrics"]["throughput"]["baseline"] == 10.0 - def test_build_latest_summary_keeps_identity_cohorts_separate(): records = [ _record( @@ -380,10 +375,11 @@ def test_partial_v2_identity_does_not_cross_model_or_gpu(): rows = build_latest_summary(records) trends = build_trends(records) - assert {(row["model_id"], row["gpu_type"]) for row in rows} == { - ("wan", "NVIDIA L40S"), - ("ltx", "NVIDIA H100"), - } + assert {(row["model_id"], row["gpu_type"]) + for row in rows} == { + ("wan", "NVIDIA L40S"), + ("ltx", "NVIDIA H100"), + } assert len(trends) == 2 @@ -509,10 +505,11 @@ def test_load_records_can_filter_baseline_eligible_records(tmp_path): records = hf_store.load_records(str(tmp_path), successful_only=True, baseline_eligible_only=True) assert len(records) == 2 - assert {record["timestamp"] for record in records} == { - "2026-01-02T00:00:00+00:00", - "2026-01-03T00:00:00+00:00", - } + assert {record["timestamp"] + for record in records} == { + "2026-01-02T00:00:00+00:00", + "2026-01-03T00:00:00+00:00", + } def test_load_records_for_model_filters_identity_cohort(tmp_path): diff --git a/fastvideo/tests/performance/test_identity.py b/fastvideo/tests/performance/test_identity.py index 105c79b24..4b7c5b8fa 100644 --- a/fastvideo/tests/performance/test_identity.py +++ b/fastvideo/tests/performance/test_identity.py @@ -99,8 +99,7 @@ def test_recipe_requires_first_class_benchmark_identity(): def test_semantically_equivalent_sequence_forms_match(): cfg = _benchmark_config() equivalent = deepcopy(cfg) - equivalent["init_kwargs"]["text_encoder_precisions"] = tuple( - equivalent["init_kwargs"]["text_encoder_precisions"]) + equivalent["init_kwargs"]["text_encoder_precisions"] = tuple(equivalent["init_kwargs"]["text_encoder_precisions"]) assert _fingerprint(cfg) == _fingerprint(equivalent) @@ -120,14 +119,14 @@ def test_runtime_resolved_values_do_not_change_recipe_fingerprint(): # degraded numbers seed a new baseline. Declared inputs (including # requested_backend) stay in the hash. cfg = _benchmark_config() - recipe = build_recipe_from_benchmark_config( - cfg, attention_backend="FLASH_ATTN", - resolved_attention_backend="FLASH_ATTN", - resolved_model_revision="aaaa1111") - changed = build_recipe_from_benchmark_config( - cfg, attention_backend="FLASH_ATTN", - resolved_attention_backend="TORCH_SDPA", - resolved_model_revision="bbbb2222") + recipe = build_recipe_from_benchmark_config(cfg, + attention_backend="FLASH_ATTN", + resolved_attention_backend="FLASH_ATTN", + resolved_model_revision="aaaa1111") + changed = build_recipe_from_benchmark_config(cfg, + attention_backend="FLASH_ATTN", + resolved_attention_backend="TORCH_SDPA", + resolved_model_revision="bbbb2222") assert recipe["model"]["resolved_revision"] == "aaaa1111" # kept for audit assert recipe["attention"]["resolved_backend"] == "FLASH_ATTN" # kept for audit @@ -167,8 +166,7 @@ def test_benchmark_version_changes_recipe_fingerprint(): def test_attention_backend_changes_recipe_fingerprint(): cfg = _benchmark_config() - assert _fingerprint(cfg, attention_backend="FLASH_ATTN") != _fingerprint( - cfg, attention_backend="TORCH_SDPA") + assert _fingerprint(cfg, attention_backend="FLASH_ATTN") != _fingerprint(cfg, attention_backend="TORCH_SDPA") def test_precision_changes_recipe_fingerprint(): @@ -289,8 +287,10 @@ def test_hardware_profile_id_uses_normalized_gpu_cohort(): ) assert profile == { - "device_type": "cuda", - "gpu_count": 2, + "device_type": + "cuda", + "gpu_count": + 2, "gpus": [ { "name": "NVIDIA L40S", @@ -303,12 +303,14 @@ def test_hardware_profile_id_uses_normalized_gpu_cohort(): "compute_capability": "8.9", }, ], - "interconnect": "none_or_partial", + "interconnect": + "none_or_partial", } assert hardware_profile_id(profile).startswith("hw-") def test_hardware_profile_pads_missing_requested_cuda_devices(monkeypatch): + class Props: name = "NVIDIA L40S" total_memory = 48 * 1024**3 @@ -349,14 +351,28 @@ def test_environment_metadata_is_separate_audit_fingerprint(): audit_a = environment_metadata( env={"IMAGE_VERSION": "py3.12-cuda13.0.0"}, package_versions={"triton": "3.4.1"}, - hardware={"device_type": "cuda", "gpu_count": 2}, - software={"python": "3.12", "pytorch": "2.12", "cuda": "13.0"}, + hardware={ + "device_type": "cuda", + "gpu_count": 2 + }, + software={ + "python": "3.12", + "pytorch": "2.12", + "cuda": "13.0" + }, ) audit_b = environment_metadata( env={"IMAGE_VERSION": "py3.12-cuda13.0.1"}, package_versions={"triton": "3.4.1"}, - hardware={"device_type": "cuda", "gpu_count": 2}, - software={"python": "3.12", "pytorch": "2.12", "cuda": "13.0"}, + hardware={ + "device_type": "cuda", + "gpu_count": 2 + }, + software={ + "python": "3.12", + "pytorch": "2.12", + "cuda": "13.0" + }, ) assert recipe_hash == _fingerprint(cfg) @@ -377,6 +393,7 @@ def test_runtime_identity_from_generator_summarizes_worker_records(): revision = "b" * 40 class FakeExecutor: + def collective_rpc(self, _method): return [ { diff --git a/fastvideo/tests/performance/test_inference_performance.py b/fastvideo/tests/performance/test_inference_performance.py index 25dec310f..99ceea70e 100644 --- a/fastvideo/tests/performance/test_inference_performance.py +++ b/fastvideo/tests/performance/test_inference_performance.py @@ -51,10 +51,8 @@ V2_REQUIRED_IDENTITY_FIELDS = ( # and the generated recipe document owns that key in emitted records. A config # declaring it now fails validation loudly instead of being silently # overwritten by the generated one. -V2_OPTIONAL_METADATA_FIELDS = ( - "quality_metadata", -) -COMMON_OBJECT_FIELDS = ("regression_thresholds",) +V2_OPTIONAL_METADATA_FIELDS = ("quality_metadata", ) +COMMON_OBJECT_FIELDS = ("regression_thresholds", ) RESULT_SCHEMA_VERSION = 2 VALID_RUN_SOURCES = {"pr", "local", "scheduled_main", "unknown"} OPTIONAL_RESULT_METADATA_FIELDS = ("quality_metadata", "variant_metadata") @@ -92,7 +90,7 @@ def _validate_integer(value, field, path): def _validate_benchmark_config(cfg, path=""): - missing_common = [field for field in ("benchmark_id",) if field not in cfg] + missing_common = [field for field in ("benchmark_id", ) if field not in cfg] if missing_common: raise ValueError(f"{path}: missing required benchmark config fields: {', '.join(missing_common)}") @@ -208,17 +206,14 @@ def _extract_component_times(result: dict) -> dict[str, float | None]: else: metric_key = STAGE_METRIC_MAP.get(stage_class) if metric_key is None: - logger.debug("Unmapped stage '%s' class '%s' (%.3fs)", - stage_name, - stage_class, + logger.debug("Unmapped stage '%s' class '%s' (%.3fs)", stage_name, stage_class, stage_data.get("execution_time", 0)) continue elapsed = stage_data.get("execution_time") if elapsed is None: continue existing = component_times[metric_key] - component_times[metric_key] = ( - elapsed if existing is None else existing + elapsed) + component_times[metric_key] = (elapsed if existing is None else existing + elapsed) return component_times @@ -260,8 +255,7 @@ def _resolve_num_gpus( ("init_kwargs.num_gpus", init_num_gpus), ("run_config.required_gpus", required_gpus), ): - if value is not None and ( - isinstance(value, bool) or not isinstance(value, int) or value < 1): + if value is not None and (isinstance(value, bool) or not isinstance(value, int) or value < 1): raise ValueError(f"{benchmark_id}: {field} must be a positive integer") parallel_sizes = [] for field in ("tp_size", "sp_size"): @@ -271,20 +265,14 @@ def _resolve_num_gpus( if isinstance(value, bool) or not isinstance(value, int) or value < 1: raise ValueError(f"{benchmark_id}: init_kwargs.{field} must be -1 or a positive integer") parallel_sizes.append(value) - if ( - init_num_gpus is not None - and required_gpus is not None - and init_num_gpus != required_gpus - ): - raise ValueError( - f"{benchmark_id}: init_kwargs.num_gpus ({init_num_gpus}) must match " - f"run_config.required_gpus ({required_gpus})") + if (init_num_gpus is not None and required_gpus is not None and init_num_gpus != required_gpus): + raise ValueError(f"{benchmark_id}: init_kwargs.num_gpus ({init_num_gpus}) must match " + f"run_config.required_gpus ({required_gpus})") declared_num_gpus = init_num_gpus or required_gpus parallel_num_gpus = max(parallel_sizes, default=1) if declared_num_gpus is not None and declared_num_gpus < parallel_num_gpus: - raise ValueError( - f"{benchmark_id}: declared GPU count ({declared_num_gpus}) must be at least " - f"max(tp_size, sp_size) ({parallel_num_gpus})") + raise ValueError(f"{benchmark_id}: declared GPU count ({declared_num_gpus}) must be at least " + f"max(tp_size, sp_size) ({parallel_num_gpus})") return declared_num_gpus or parallel_num_gpus @@ -319,7 +307,7 @@ def _collect_worker_identity(worker) -> dict[str, Any]: if isinstance(modules, Mapping): module_iter = modules.values() elif isinstance(pipeline, torch.nn.Module): - module_iter = (pipeline,) + module_iter = (pipeline, ) else: module_iter = () @@ -368,10 +356,7 @@ def _runtime_identity_from_generator(generator) -> dict[str, Any]: def _benchmark_identity_fields(cfg): identity = benchmark_identity_from_config(cfg) - return { - key: identity[key] - for key in ("workload_id", "variant_id", "benchmark_version") - } + return {key: identity[key] for key in ("workload_id", "variant_id", "benchmark_version")} def _build_identity_fields(cfg, init_kwargs, prompt, runtime_identity): @@ -457,11 +442,7 @@ def _ci_provenance_fields() -> dict[str, str]: def _configured_result_metadata(cfg: Mapping[str, Any]) -> dict[str, Any]: - return { - field: cfg[field] - for field in OPTIONAL_RESULT_METADATA_FIELDS - if field in cfg and cfg[field] is not None - } + return {field: cfg[field] for field in OPTIONAL_RESULT_METADATA_FIELDS if field in cfg and cfg[field] is not None} def _build_result_record( @@ -491,43 +472,53 @@ def _build_result_record( if isinstance(num_frames, (int, float)) and avg_time > 0: throughput_fps = num_frames / avg_time - result_schema_fields = ( - {"result_schema_version": RESULT_SCHEMA_VERSION} - if _is_v2_config(cfg) - else {} - ) + result_schema_fields = ({"result_schema_version": RESULT_SCHEMA_VERSION} if _is_v2_config(cfg) else {}) return { - "benchmark_id": cfg["benchmark_id"], + "benchmark_id": + cfg["benchmark_id"], **result_schema_fields, - "model_short_name": model_info.get("model_short_name", ""), - "device": device_name, - "num_gpus": num_gpus, - "num_warmup_runs": num_warmup, - "num_measurement_runs": num_measure, - "avg_generation_time_s": round(avg_time, 3), + "model_short_name": + model_info.get("model_short_name", ""), + "device": + device_name, + "num_gpus": + num_gpus, + "num_warmup_runs": + num_warmup, + "num_measurement_runs": + num_measure, + "avg_generation_time_s": + round(avg_time, 3), "individual_times_s": [round(t, 3) for t in times], - "throughput_fps": round(throughput_fps, 3) - if throughput_fps is not None else None, - "max_peak_memory_mb": round(max_peak_memory, 1), + "throughput_fps": + round(throughput_fps, 3) if throughput_fps is not None else None, + "max_peak_memory_mb": + round(max_peak_memory, 1), "individual_peak_memories_mb": [round(m, 1) for m in peak_memories], - "thresholds": dict(thresholds), - "regression_thresholds": cfg.get("regression_thresholds", {}), - "commit": os.environ.get("BUILDKITE_COMMIT", ""), + "thresholds": + dict(thresholds), + "regression_thresholds": + cfg.get("regression_thresholds", {}), + "commit": + os.environ.get("BUILDKITE_COMMIT", ""), **_ci_provenance_fields(), - "timestamp": timestamp or datetime.now(timezone.utc).isoformat(), + "timestamp": + timestamp or datetime.now(timezone.utc).isoformat(), **_configured_result_metadata(cfg), - "text_encoder_time_s": _avg_component(all_component_times, - "text_encoder_time_s"), - "dit_time_s": _avg_component(all_component_times, "dit_time_s"), - "vae_decode_time_s": _avg_component(all_component_times, - "vae_decode_time_s"), + "text_encoder_time_s": + _avg_component(all_component_times, "text_encoder_time_s"), + "dit_time_s": + _avg_component(all_component_times, "dit_time_s"), + "vae_decode_time_s": + _avg_component(all_component_times, "vae_decode_time_s"), **_build_identity_fields(cfg, init_kwargs, prompt, runtime_identity), } # -- Test ------------------------------------------------------------------- + def _run_benchmark(cfg): run_config = cfg.get("run_config") or {} model_info = cfg["model"] @@ -553,8 +544,7 @@ def _run_benchmark(cfg): # Output directory for generated videos script_dir = os.path.dirname(os.path.abspath(__file__)) - output_dir = os.path.join(script_dir, "generated_videos", - cfg["benchmark_id"]) + output_dir = os.path.join(script_dir, "generated_videos", cfg["benchmark_id"]) os.makedirs(output_dir, exist_ok=True) gen_kwargs["output_path"] = output_dir @@ -606,21 +596,18 @@ def _run_benchmark(cfg): device_name=device_name, ) - logger.info( - "Performance results: avg_time=%.2fs, " - "max_peak_memory=%.0fMB", avg_time, max_peak_memory) + logger.info("Performance results: avg_time=%.2fs, " + "max_peak_memory=%.0fMB", avg_time, max_peak_memory) _write_results(results) max_time = thresholds["max_generation_time_s"] max_mem = thresholds["max_peak_memory_mb"] - assert avg_time <= max_time, ( - f"Average generation time {avg_time:.2f}s exceeds " - f"threshold {max_time:.1f}s for {device_name}") + assert avg_time <= max_time, (f"Average generation time {avg_time:.2f}s exceeds " + f"threshold {max_time:.1f}s for {device_name}") - assert max_peak_memory <= max_mem, ( - f"Peak memory {max_peak_memory:.0f}MB exceeds " - f"threshold {max_mem:.0f}MB for {device_name}") + assert max_peak_memory <= max_mem, (f"Peak memory {max_peak_memory:.0f}MB exceeds " + f"threshold {max_mem:.0f}MB for {device_name}") component_thresholds = { "text_encoder_time_s": thresholds.get("max_text_encoder_time_s"), @@ -632,9 +619,8 @@ def _run_benchmark(cfg): continue actual = results[metric] if actual is not None: - assert actual <= max_val, ( - f"{metric} {actual:.3f}s exceeds threshold {max_val:.3f}s " - f"for {device_name}") + assert actual <= max_val, (f"{metric} {actual:.3f}s exceeds threshold {max_val:.3f}s " + f"for {device_name}") @pytest.mark.parametrize( diff --git a/fastvideo/tests/performance/test_inference_performance_component_times.py b/fastvideo/tests/performance/test_inference_performance_component_times.py index ff2512c8c..209813f26 100644 --- a/fastvideo/tests/performance/test_inference_performance_component_times.py +++ b/fastvideo/tests/performance/test_inference_performance_component_times.py @@ -87,9 +87,15 @@ def test_extract_component_times_keeps_legacy_class_name_keys(): result = { "logging_info": { "stages": { - "TextEncodingStage": {"execution_time": 1.0}, - "DenoisingStage": {"execution_time": 2.0}, - "DecodingStage": {"execution_time": 3.0}, + "TextEncodingStage": { + "execution_time": 1.0 + }, + "DenoisingStage": { + "execution_time": 2.0 + }, + "DecodingStage": { + "execution_time": 3.0 + }, }, }, } @@ -130,9 +136,15 @@ def test_extract_component_times_ignores_unmapped_stages(): result = { "logging_info": { "stages": { - "PostDecodeFrameProcessStage": {"execution_time": 0.2}, - "VideoSaveStage": {"execution_time": 0.4}, - "AudioMuxStage": {"execution_time": 0.1}, + "PostDecodeFrameProcessStage": { + "execution_time": 0.2 + }, + "VideoSaveStage": { + "execution_time": 0.4 + }, + "AudioMuxStage": { + "execution_time": 0.1 + }, }, }, } diff --git a/fastvideo/tests/performance/test_inference_performance_identity.py b/fastvideo/tests/performance/test_inference_performance_identity.py index 8391e56f2..f8eda3340 100644 --- a/fastvideo/tests/performance/test_inference_performance_identity.py +++ b/fastvideo/tests/performance/test_inference_performance_identity.py @@ -201,9 +201,7 @@ def test_producer_tracks_runtime_software_identity_and_container_audit(monkeypat } assert identity_fields["software_profile"] == expected_profile assert identity_fields["software_profile_id"] == perf.software_profile_id(expected_profile) - assert identity_fields["environment_metadata"]["env"]["FASTVIDEO_CONTAINER_IMAGE_REF"].endswith( - "sha256:abc" - ) + assert identity_fields["environment_metadata"]["env"]["FASTVIDEO_CONTAINER_IMAGE_REF"].endswith("sha256:abc") first_profile_id = identity_fields["software_profile_id"] monkeypatch.setenv( @@ -220,9 +218,7 @@ def test_producer_tracks_runtime_software_identity_and_container_audit(monkeypat }, ) assert changed_audit["software_profile_id"] == first_profile_id - assert changed_audit["environment_metadata"]["env"]["FASTVIDEO_CONTAINER_IMAGE_REF"].endswith( - "sha256:def" - ) + assert changed_audit["environment_metadata"]["env"]["FASTVIDEO_CONTAINER_IMAGE_REF"].endswith("sha256:def") monkeypatch.setenv("FASTVIDEO_FA4", "0") changed_runtime = perf._build_identity_fields( diff --git a/fastvideo/tests/performance/test_inference_performance_result_schema.py b/fastvideo/tests/performance/test_inference_performance_result_schema.py index c09b6ef9d..30c6191a8 100644 --- a/fastvideo/tests/performance/test_inference_performance_result_schema.py +++ b/fastvideo/tests/performance/test_inference_performance_result_schema.py @@ -177,7 +177,10 @@ def test_resolve_num_gpus_uses_one_count_and_rejects_conflicts(): with pytest.raises(ValueError, match=r"at least max\(tp_size, sp_size\)"): perf_test._resolve_num_gpus( - {"num_gpus": 1, "tp_size": 2}, + { + "num_gpus": 1, + "tp_size": 2 + }, {"required_gpus": 1}, "wan-t2v-1.3b-tp2", ) diff --git a/fastvideo/tests/ssim/bootstrap_references.py b/fastvideo/tests/ssim/bootstrap_references.py index c0ac388d7..e9c380369 100644 --- a/fastvideo/tests/ssim/bootstrap_references.py +++ b/fastvideo/tests/ssim/bootstrap_references.py @@ -36,8 +36,7 @@ def xfail_missing_reference_in_bootstrap_mode( generated_path = Path(generated_artifact_path) if not generated_path.exists(): raise FileNotFoundError( - f"SSIM bootstrap mode is enabled, but generated {artifact_kind} artifact is missing: {generated_path}" - ) + f"SSIM bootstrap mode is enabled, but generated {artifact_kind} artifact is missing: {generated_path}") repo_id = os.environ.get(HF_REPO_ENV_KEY, DEFAULT_REPO_ID) repo_type = os.environ.get(HF_REPO_TYPE_ENV_KEY, DEFAULT_REPO_TYPE) @@ -47,10 +46,8 @@ def xfail_missing_reference_in_bootstrap_mode( generated_artifact_path=generated_path, reference_folder=Path(reference_folder), ) - pytest.xfail( - "SSIM bootstrap mode generated a draft " - f"{artifact_kind} reference at {repo_id}/{draft_path}. " - "Review it, then promote with " - "`python fastvideo/tests/ssim/reference_videos_cli.py promote-draft " - f"--quality-tier {get_output_quality_tier()} --model-id `." - ) + pytest.xfail("SSIM bootstrap mode generated a draft " + f"{artifact_kind} reference at {repo_id}/{draft_path}. " + "Review it, then promote with " + "`python fastvideo/tests/ssim/reference_videos_cli.py promote-draft " + f"--quality-tier {get_output_quality_tier()} --model-id `.") diff --git a/fastvideo/tests/ssim/conftest.py b/fastvideo/tests/ssim/conftest.py index b76447106..d93a59eb1 100644 --- a/fastvideo/tests/ssim/conftest.py +++ b/fastvideo/tests/ssim/conftest.py @@ -16,18 +16,14 @@ def pytest_addoption(parser): "--ssim-full-quality", action="store_true", default=False, - help=( - "Use *_FULL_QUALITY_PARAMS for SSIM tests. " - "Default keeps the original CI-friendly params." - ), + help=("Use *_FULL_QUALITY_PARAMS for SSIM tests. " + "Default keeps the original CI-friendly params."), ) parser.addoption( "--ssim-reference-repo", default=os.environ.get(HF_REPO_ENV_KEY, ""), - help=( - "HF repo id for SSIM reference videos " - f"(overrides {HF_REPO_ENV_KEY})." - ), + help=("HF repo id for SSIM reference videos " + f"(overrides {HF_REPO_ENV_KEY})."), ) parser.addoption( "--skip-ssim-reference-download", @@ -55,9 +51,8 @@ def pytest_configure(config): os.environ[BOOTSTRAP_ENV_KEY] = "1" skip_download = config.getoption("--skip-ssim-reference-download") - skip_download = skip_download or os.environ.get( - "FASTVIDEO_SSIM_SKIP_REFERENCE_DOWNLOAD", "" - ).strip().lower() in {"1", "true", "yes", "on"} + skip_download = skip_download or os.environ.get("FASTVIDEO_SSIM_SKIP_REFERENCE_DOWNLOAD", + "").strip().lower() in {"1", "true", "yes", "on"} if not skip_download: ensure_reference_videos_available( diff --git a/fastvideo/tests/ssim/latent_similarity_utils.py b/fastvideo/tests/ssim/latent_similarity_utils.py index e2a06d8e3..9dbb15670 100644 --- a/fastvideo/tests/ssim/latent_similarity_utils.py +++ b/fastvideo/tests/ssim/latent_similarity_utils.py @@ -64,8 +64,7 @@ from torch.nn.functional import cosine_similarity from fastvideo import VideoGenerator from fastvideo.tests.ssim.bootstrap_references import ( - xfail_missing_reference_in_bootstrap_mode, -) + xfail_missing_reference_in_bootstrap_mode, ) from fastvideo.tests.ssim.inference_similarity_utils import ( attention_backend, build_generation_kwargs, @@ -108,34 +107,22 @@ def _extract_expected_slice( kind = spec.get("kind", "corner_3x3_first_frame") if kind == "corner_3x3_first_frame": if latent.dim() != 5: - raise ValueError( - f"corner_3x3_first_frame requires 5-D latent [B,C,T,H,W]; " - f"got shape {tuple(latent.shape)}") + raise ValueError(f"corner_3x3_first_frame requires 5-D latent [B,C,T,H,W]; " + f"got shape {tuple(latent.shape)}") _, _, t, h, w = latent.shape if t < 1 or h < 3 or w < 3: - raise ValueError( - "corner_3x3_first_frame requires T>=1, H>=3, W>=3; got " - f"shape {tuple(latent.shape)}") - return (latent[0, :, 0, :3, :3] - .detach() - .to(torch.float32) - .reshape(-1) - .contiguous()) + raise ValueError("corner_3x3_first_frame requires T>=1, H>=3, W>=3; got " + f"shape {tuple(latent.shape)}") + return (latent[0, :, 0, :3, :3].detach().to(torch.float32).reshape(-1).contiguous()) if kind == "audio_first_8_timesteps": if latent.dim() != 3: - raise ValueError( - f"audio_first_8_timesteps requires 3-D latent [B,C,T]; " - f"got shape {tuple(latent.shape)}") + raise ValueError(f"audio_first_8_timesteps requires 3-D latent [B,C,T]; " + f"got shape {tuple(latent.shape)}") _, _, t = latent.shape if t < 8: - raise ValueError( - "audio_first_8_timesteps requires T>=8; got " - f"shape {tuple(latent.shape)}") - return (latent[0, :, :8] - .detach() - .to(torch.float32) - .reshape(-1) - .contiguous()) + raise ValueError("audio_first_8_timesteps requires T>=8; got " + f"shape {tuple(latent.shape)}") + return (latent[0, :, :8].detach().to(torch.float32).reshape(-1).contiguous()) raise ValueError(f"Unknown slice kind: {kind!r}") @@ -144,9 +131,7 @@ def _cosine_distance(a: torch.Tensor, b: torch.Tensor) -> float: a32 = a.detach().to(torch.float32).reshape(-1) b32 = b.detach().to(torch.float32).reshape(-1) if a32.shape != b32.shape: - raise ValueError( - f"Cosine shape mismatch: {tuple(a32.shape)} vs {tuple(b32.shape)}" - ) + raise ValueError(f"Cosine shape mismatch: {tuple(a32.shape)} vs {tuple(b32.shape)}") sim = cosine_similarity(a32.unsqueeze(0), b32.unsqueeze(0), dim=1).item() return 1.0 - float(sim) @@ -208,11 +193,10 @@ def load_latent_reference(path: str) -> dict[str, Any]: payload = torch.load(path, map_location="cpu", weights_only=False) fmt = payload.get("format_version") if isinstance(payload, dict) else None if fmt != LATENT_REFERENCE_FORMAT_VERSION: - raise ValueError( - f"Latent reference at {path!r} has format_version={fmt!r}; " - f"expected {LATENT_REFERENCE_FORMAT_VERSION}. Re-seed via the " - "test that produces this artefact, then re-upload through " - "fastvideo/tests/ssim/reference_videos_cli.py upload.") + raise ValueError(f"Latent reference at {path!r} has format_version={fmt!r}; " + f"expected {LATENT_REFERENCE_FORMAT_VERSION}. Re-seed via the " + "test that produces this artefact, then re-upload through " + "fastvideo/tests/ssim/reference_videos_cli.py upload.") return payload @@ -287,10 +271,9 @@ def _assert_latent_similarity( slice_spec = ref.get("slice_spec", DEFAULT_SLICE_SPEC) if tuple(gen_latent.shape) != ref_shape: - raise AssertionError( - f"Generated latent shape {tuple(gen_latent.shape)} does not " - f"match reference shape {ref_shape} for {model_id} with " - f"backend {attention_backend_name}") + raise AssertionError(f"Generated latent shape {tuple(gen_latent.shape)} does not " + f"match reference shape {ref_shape} for {model_id} with " + f"backend {attention_backend_name}") gen_slice = _extract_expected_slice(gen_latent, slice_spec) slice_cos = _cosine_distance(gen_slice, expected_slice) @@ -301,8 +284,7 @@ def _assert_latent_similarity( # fp32 (see :func:`save_latent_reference`). gen_full_matched = gen_latent.to(torch.float16).to(torch.float32) full_cos = _cosine_distance(gen_full_matched, ref_full) - max_abs_diff = float( - (gen_full_matched - ref_full).abs().max().item()) + max_abs_diff = float((gen_full_matched - ref_full).abs().max().item()) metrics: dict[str, float] = { "slice_cosine_distance": slice_cos, @@ -318,13 +300,11 @@ def _assert_latent_similarity( failures: list[str] = [] if slice_cos > slice_cosine_threshold: - failures.append( - f"slice cosine {slice_cos:.6e} > threshold " - f"{slice_cosine_threshold:.6e}") + failures.append(f"slice cosine {slice_cos:.6e} > threshold " + f"{slice_cosine_threshold:.6e}") if full_cos > full_cosine_threshold: - failures.append( - f"full cosine {full_cos:.6e} > threshold " - f"{full_cosine_threshold:.6e}") + failures.append(f"full cosine {full_cos:.6e} > threshold " + f"{full_cosine_threshold:.6e}") passed = not failures write_latent_similarity_results( @@ -343,10 +323,9 @@ def _assert_latent_similarity( ) if failures: - raise AssertionError( - f"Latent regression exceeded tolerance for {model_id} with " - f"backend {attention_backend_name}: {'; '.join(failures)}. " - f"Full metrics: {metrics}") + raise AssertionError(f"Latent regression exceeded tolerance for {model_id} with " + f"backend {attention_backend_name}: {'; '.join(failures)}. " + f"Full metrics: {metrics}") return metrics @@ -356,22 +335,18 @@ def _extract_latent_from_result(result: Any) -> torch.Tensor: ``generate_video`` output. """ if not isinstance(result, dict): - raise RuntimeError( - "VideoGenerator.generate_video returned unexpected payload " - f"(type={type(result)!r}); expected dict with 'samples'.") + raise RuntimeError("VideoGenerator.generate_video returned unexpected payload " + f"(type={type(result)!r}); expected dict with 'samples'.") samples = result.get("samples") if samples is None: - raise RuntimeError( - "VideoGenerator did not return latent samples. Ensure " - "output_type='latent' and return_frames=True for this call.") + raise RuntimeError("VideoGenerator did not return latent samples. Ensure " + "output_type='latent' and return_frames=True for this call.") if not isinstance(samples, torch.Tensor): - raise RuntimeError( - f"Expected torch.Tensor samples; got type={type(samples)!r}.") + raise RuntimeError(f"Expected torch.Tensor samples; got type={type(samples)!r}.") gen_latent = samples.detach().to(torch.float32).cpu() if gen_latent.dim() not in (3, 5): - raise RuntimeError( - "Expected 5-D video latent (B,C,T,H,W) or 3-D audio latent " - f"(B,C,T); got shape {tuple(gen_latent.shape)}") + raise RuntimeError("Expected 5-D video latent (B,C,T,H,W) or 3-D audio latent " + f"(B,C,T); got shape {tuple(gen_latent.shape)}") return gen_latent @@ -473,10 +448,9 @@ def run_text_to_latent_similarity_test( reference_folder=reference_folder, artifact_kind="latent", ) - raise FileNotFoundError( - f"Reference folder does not exist: {reference_folder}\n" - f"To download references, run:\n" - f" python fastvideo/tests/ssim/reference_videos_cli.py download") + raise FileNotFoundError(f"Reference folder does not exist: {reference_folder}\n" + f"To download references, run:\n" + f" python fastvideo/tests/ssim/reference_videos_cli.py download") reference_latent_path = os.path.join( reference_folder, @@ -488,9 +462,8 @@ def run_text_to_latent_similarity_test( reference_folder=reference_folder, artifact_kind="latent", ) - raise FileNotFoundError( - "Reference latent missing for prompt/backend: " - f"{reference_latent_path}") + raise FileNotFoundError("Reference latent missing for prompt/backend: " + f"{reference_latent_path}") return _assert_latent_similarity( logger=logger, diff --git a/fastvideo/tests/ssim/reference_utils.py b/fastvideo/tests/ssim/reference_utils.py index 78ee60b2f..e7390fadb 100644 --- a/fastvideo/tests/ssim/reference_utils.py +++ b/fastvideo/tests/ssim/reference_utils.py @@ -89,9 +89,7 @@ def build_reference_folder_path( model_id, attention_backend, ) - if quality_tier == DEFAULT_OUTPUT_QUALITY_TIER and os.path.exists( - legacy_reference_folder - ): + if quality_tier == DEFAULT_OUTPUT_QUALITY_TIER and os.path.exists(legacy_reference_folder): return legacy_reference_folder return tiered_reference_folder diff --git a/fastvideo/tests/ssim/reference_videos_cli.py b/fastvideo/tests/ssim/reference_videos_cli.py index 8b9c99ec6..4c71c46f2 100644 --- a/fastvideo/tests/ssim/reference_videos_cli.py +++ b/fastvideo/tests/ssim/reference_videos_cli.py @@ -11,14 +11,13 @@ import tempfile from collections.abc import Iterable, Sequence from pathlib import Path - VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv") # Additional artefact types stored under the same reference folders. Latent # tensors (`.pt`) back the latent-slice regression tests used by flaky # pixel-space models (e.g. LTX-2 distilled). Keeping them in the same upload # flow means seeding a new test only requires one HF round-trip. -LATENT_EXTENSIONS = (".pt",) -REFERENCE_EXTENSIONS = VIDEO_EXTENSIONS + LATENT_EXTENSIONS + (".png",) +LATENT_EXTENSIONS = (".pt", ) +REFERENCE_EXTENSIONS = VIDEO_EXTENSIONS + LATENT_EXTENSIONS + (".png", ) HF_TOKEN_ENV_KEYS = ("HF_API_KEY", "HUGGINGFACE_HUB_TOKEN", "HF_TOKEN") HF_REPO_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO" HF_REPO_TYPE_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE" @@ -171,8 +170,7 @@ def _load_hf_sdk(): from huggingface_hub import HfApi, snapshot_download except ImportError as exc: raise RuntimeError( - "huggingface_hub is required for download/upload.\nInstall with: uv pip install huggingface_hub" - ) from exc + "huggingface_hub is required for download/upload.\nInstall with: uv pip install huggingface_hub") from exc return HfApi, snapshot_download @@ -249,8 +247,7 @@ def download_reference_videos( ) if copied_roots == 0: raise RuntimeError( - "HF download completed but no reference_videos content was found in downloaded artifacts." - ) + "HF download completed but no reference_videos content was found in downloaded artifacts.") for quality_tier in tiers: marker_path = local_dir / REFERENCE_VIDEOS_DIRNAME / quality_tier / f".download_complete_{quality_tier}" marker_path.parent.mkdir(parents=True, exist_ok=True) @@ -277,8 +274,7 @@ def upload_reference_videos( ) try: - existing_repo_files = set( - api.list_repo_files(repo_id=repo_id, repo_type=repo_type)) + existing_repo_files = set(api.list_repo_files(repo_id=repo_id, repo_type=repo_type)) except Exception: # Fresh repo or list failure — treat as empty so upload can proceed. existing_repo_files = set() @@ -290,23 +286,19 @@ def upload_reference_videos( if model_id: folder_path = reference_dir / model_id if not folder_path.exists(): - raise FileNotFoundError( - f"Model subfolder not found for upload: {folder_path}") + raise FileNotFoundError(f"Model subfolder not found for upload: {folder_path}") path_in_repo = f"{base_in_repo}/{model_id}" else: folder_path = reference_dir path_in_repo = base_in_repo - conflicts = sorted( - f for f in existing_repo_files - if f.startswith(f"{path_in_repo}/") or f == path_in_repo) + conflicts = sorted(f for f in existing_repo_files if f.startswith(f"{path_in_repo}/") or f == path_in_repo) if conflicts and not force: preview = "\n".join(f" - {c}" for c in conflicts[:10]) more = f"\n ... and {len(conflicts) - 10} more" if len(conflicts) > 10 else "" - raise RuntimeError( - f"Refusing to overwrite existing HF files under {path_in_repo} " - f"({len(conflicts)} file(s) already present):\n{preview}{more}\n" - f"Re-run with --force to overwrite.") + raise RuntimeError(f"Refusing to overwrite existing HF files under {path_in_repo} " + f"({len(conflicts)} file(s) already present):\n{preview}{more}\n" + f"Re-run with --force to overwrite.") target_desc = f"{reference_dir.name}/{model_id}" if model_id else reference_dir.name print(f"Uploading {target_desc} ({quality_tier}) to {repo_id}/{path_in_repo} ...") @@ -330,8 +322,7 @@ def _reference_folder_repo_relative(reference_folder: Path, base_dir: Path | Non legacy_relative = resolved_reference_folder.relative_to(resolved_base_dir) except ValueError as exc: raise RuntimeError( - f"Reference folder {reference_folder} is not under SSIM directory {resolved_base_dir}" - ) from exc + f"Reference folder {reference_folder} is not under SSIM directory {resolved_base_dir}") from exc if not legacy_relative.parts or not legacy_relative.parts[0].endswith("_reference_videos"): raise RuntimeError( f"Reference folder {reference_folder} does not match reference_videos////" @@ -339,9 +330,7 @@ def _reference_folder_repo_relative(reference_folder: Path, base_dir: Path | Non relative = Path(DEFAULT_OUTPUT_QUALITY_TIER, *legacy_relative.parts) if len(relative.parts) < 4: - raise RuntimeError( - f"Reference folder {reference_folder} does not include ///" - ) + raise RuntimeError(f"Reference folder {reference_folder} does not include ///") return relative @@ -357,10 +346,8 @@ def upload_draft_reference_artifact( ) -> str: resolved_token = token or _get_hf_token() if resolved_token is None: - raise RuntimeError( - "Hugging Face API key is required for SSIM bootstrap draft upload. " - "Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN." - ) + raise RuntimeError("Hugging Face API key is required for SSIM bootstrap draft upload. " + "Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN.") if not generated_artifact_path.exists(): raise FileNotFoundError(f"Generated artifact not found: {generated_artifact_path}") @@ -438,16 +425,13 @@ def promote_draft_references( if not draft_files: raise RuntimeError(f"No draft references found under {repo_id}/{draft_prefix}") - conflicts = sorted( - f for f in existing_repo_files - if f.startswith(f"{reference_prefix}/") or f == reference_prefix) + conflicts = sorted(f for f in existing_repo_files if f.startswith(f"{reference_prefix}/") or f == reference_prefix) if conflicts and not force: preview = "\n".join(f" - {c}" for c in conflicts[:10]) more = f"\n ... and {len(conflicts) - 10} more" if len(conflicts) > 10 else "" - raise RuntimeError( - f"Refusing to overwrite existing HF files under {reference_prefix} " - f"({len(conflicts)} file(s) already present):\n{preview}{more}\n" - f"Re-run with --force to overwrite.") + raise RuntimeError(f"Refusing to overwrite existing HF files under {reference_prefix} " + f"({len(conflicts)} file(s) already present):\n{preview}{more}\n" + f"Re-run with --force to overwrite.") with tempfile.TemporaryDirectory(prefix="fv2-ssim-draft-") as tmp_cache_dir: snapshot_path = Path( @@ -457,8 +441,7 @@ def promote_draft_references( cache_dir=tmp_cache_dir, token=token, allow_patterns=[f"{draft_prefix}/**"], - ) - ) + )) draft_root = snapshot_path / draft_prefix if not draft_root.exists(): raise RuntimeError(f"Downloaded draft path is missing: {draft_root}") @@ -518,8 +501,7 @@ def ensure_reference_videos_available( resolved_repo_type = repo_type or _default_repo_type() if not resolved_repo_id: raise RuntimeError( - f"No local reference videos found and no HF repo configured.\nSet {HF_REPO_ENV_KEY} or pass --repo-id." - ) + f"No local reference videos found and no HF repo configured.\nSet {HF_REPO_ENV_KEY} or pass --repo-id.") print(f"Repo ID: {resolved_repo_id}") print(f"Quality tier: {quality_tier}") @@ -534,11 +516,9 @@ def ensure_reference_videos_available( print(f"Download completed for {quality_tier} reference videos.") except Exception as exc: print(f"ERROR: Failed to download {quality_tier} reference videos from {resolved_repo_id}.") - print( - f"Suggested command to retry: " - f"python fastvideo/tests/ssim/reference_videos_cli.py download " - f"--quality-tier {quality_tier}" - ) + print(f"Suggested command to retry: " + f"python fastvideo/tests/ssim/reference_videos_cli.py download " + f"--quality-tier {quality_tier}") raise if not _has_local_reference_videos(target_dir, quality_tier): @@ -549,25 +529,21 @@ def ensure_reference_videos_available( def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( prog="reference_videos_cli", - description=( - "Manage SSIM reference videos locally and on Hugging Face.\n" - "This tool can copy generated videos into reference folders,\n" - "download references from a public HF repo, and upload updates." - ), - epilog=( - "Examples:\n" - " # Copy default generated videos into local L40S refs\n" - " python fastvideo/tests/ssim/reference_videos_cli.py copy-local \\\n" - " --quality-tier default \\\n" - " --device-folder L40S_reference_videos \\\n" - " --reference-dir fastvideo/tests/ssim/reference_videos/default/L40S_reference_videos\n\n" - " # Download both default and full-quality references from HF\n" - " python fastvideo/tests/ssim/reference_videos_cli.py download \\\n" - " --repo-id FastVideo/ssim-reference-videos\n\n" - " # Upload both tiers and all GPU folders (fails if token is unset)\n" - " python fastvideo/tests/ssim/reference_videos_cli.py upload \\\n" - " --repo-id FastVideo/ssim-reference-videos" - ), + description=("Manage SSIM reference videos locally and on Hugging Face.\n" + "This tool can copy generated videos into reference folders,\n" + "download references from a public HF repo, and upload updates."), + epilog=("Examples:\n" + " # Copy default generated videos into local L40S refs\n" + " python fastvideo/tests/ssim/reference_videos_cli.py copy-local \\\n" + " --quality-tier default \\\n" + " --device-folder L40S_reference_videos \\\n" + " --reference-dir fastvideo/tests/ssim/reference_videos/default/L40S_reference_videos\n\n" + " # Download both default and full-quality references from HF\n" + " python fastvideo/tests/ssim/reference_videos_cli.py download \\\n" + " --repo-id FastVideo/ssim-reference-videos\n\n" + " # Upload both tiers and all GPU folders (fails if token is unset)\n" + " python fastvideo/tests/ssim/reference_videos_cli.py upload \\\n" + " --repo-id FastVideo/ssim-reference-videos"), formatter_class=argparse.RawTextHelpFormatter, ) @@ -624,10 +600,8 @@ def _build_parser() -> argparse.ArgumentParser: "--local-dir", type=Path, default=_ssim_dir(), - help=( - "Local SSIM directory where references are stored under " - "reference_videos//_reference_videos." - ), + help=("Local SSIM directory where references are stored under " + "reference_videos//_reference_videos."), ) download_parser.add_argument( "--quality-tier", @@ -664,11 +638,9 @@ def _build_parser() -> argparse.ArgumentParser: "--reference-dir", action="append", default=[], - help=( - "Reference directory to upload (repeatable). " - "Requires single --quality-tier. " - "Default: discover folders under --base-dir for chosen tier(s)." - ), + help=("Reference directory to upload (repeatable). " + "Requires single --quality-tier. " + "Default: discover folders under --base-dir for chosen tier(s)."), ) upload_parser.add_argument( "--quality-tier", @@ -700,18 +672,16 @@ def _build_parser() -> argparse.ArgumentParser: upload_parser.add_argument( "--model-id", default=None, - help=( - "Restrict upload to a single model subfolder " - "(reference_videos///). " - "Use when seeding references for a single new test."), + help=("Restrict upload to a single model subfolder " + "(reference_videos///). " + "Use when seeding references for a single new test."), ) upload_parser.add_argument( "--force", action="store_true", - help=( - "Allow overwriting files that already exist at the target path on " - "Hugging Face. Off by default so seeding a new test cannot " - "clobber existing references."), + help=("Allow overwriting files that already exist at the target path on " + "Hugging Face. Off by default so seeding a new test cannot " + "clobber existing references."), ) promote_parser = subparsers.add_parser( @@ -814,11 +784,8 @@ def main(argv: Sequence[str] | None = None) -> int: else: # Backward compatibility for pre-device-folder layout. generated_dir = tier_root - reference_dir = ( - args.reference_dir - if args.reference_dir is not None - else (_ssim_dir() / REFERENCE_VIDEOS_DIRNAME / args.quality_tier / device_folder) - ) + reference_dir = (args.reference_dir if args.reference_dir is not None else + (_ssim_dir() / REFERENCE_VIDEOS_DIRNAME / args.quality_tier / device_folder)) copied = copy_generated_to_reference( generated_dir=generated_dir, reference_dir=reference_dir, @@ -826,9 +793,8 @@ def main(argv: Sequence[str] | None = None) -> int: ) print(f"Done. {'Would copy' if args.dry_run else 'Copied'} {copied} video files.") if not args.dry_run and copied > 0: - marker = ( - _ssim_dir() / REFERENCE_VIDEOS_DIRNAME / args.quality_tier / f".download_complete_{args.quality_tier}" - ) + marker = (_ssim_dir() / REFERENCE_VIDEOS_DIRNAME / args.quality_tier / + f".download_complete_{args.quality_tier}") marker.parent.mkdir(parents=True, exist_ok=True) marker.touch() return 0 @@ -849,8 +815,7 @@ def main(argv: Sequence[str] | None = None) -> int: token = _get_hf_token() if token is None: raise RuntimeError( - "Hugging Face API key is required for upload. Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN." - ) + "Hugging Face API key is required for upload. Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN.") reference_dirs_by_tier = _resolve_upload_reference_dirs( base_dir=args.base_dir.resolve(), @@ -860,8 +825,7 @@ def main(argv: Sequence[str] | None = None) -> int: ) if not reference_dirs_by_tier: raise RuntimeError( - f"No *_reference_videos directories found for the selected quality tier(s) under {args.base_dir}" - ) + f"No *_reference_videos directories found for the selected quality tier(s) under {args.base_dir}") upload_reference_videos( repo_id=args.repo_id, @@ -878,10 +842,8 @@ def main(argv: Sequence[str] | None = None) -> int: if args.command == "promote-draft": token = _get_hf_token() if token is None: - raise RuntimeError( - "Hugging Face API key is required for promote-draft. " - "Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN." - ) + raise RuntimeError("Hugging Face API key is required for promote-draft. " + "Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN.") promote_draft_references( repo_id=args.repo_id, repo_type=args.repo_type, diff --git a/fastvideo/tests/ssim/test_causal_similarity.py b/fastvideo/tests/ssim/test_causal_similarity.py index ed04bf3fd..27dd8fc61 100644 --- a/fastvideo/tests/ssim/test_causal_similarity.py +++ b/fastvideo/tests/ssim/test_causal_similarity.py @@ -28,8 +28,7 @@ SF_WAN_T2V_PARAMS = { "tp_size": 1, } -_SF_WAN_T2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers") +_SF_WAN_T2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers") SF_WAN_T2V_FULL_QUALITY_PARAMS = { "num_gpus": SF_WAN_T2V_PARAMS["num_gpus"], "model_path": SF_WAN_T2V_PARAMS["model_path"], diff --git a/fastvideo/tests/ssim/test_dreamx_world_similarity.py b/fastvideo/tests/ssim/test_dreamx_world_similarity.py index 1205ca804..c4ef43f0b 100644 --- a/fastvideo/tests/ssim/test_dreamx_world_similarity.py +++ b/fastvideo/tests/ssim/test_dreamx_world_similarity.py @@ -50,7 +50,6 @@ DREAMX_WORLD_FULL_QUALITY_PARAMS = { "guidance_scale": 5.0, } - DREAMX_WORLD_AR_PARAMS = { "num_gpus": 1, "model_path": _AR_MODEL_PATH, @@ -95,7 +94,6 @@ DREAMX_WORLD_TEST_CASES = [ ), ] - DREAMX_WORLD_AR_TEST_CASES = [ ( "A long autonomous drive through a futuristic coastal city at sunrise, " @@ -156,6 +154,7 @@ def test_dreamx_world_inference_similarity( }, ) + @pytest.mark.parametrize(("prompt", "action_list", "action_speed_list"), DREAMX_WORLD_AR_TEST_CASES) @pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"]) @pytest.mark.parametrize("model_id", list(DREAMX_WORLD_AR_MODEL_TO_PARAMS.keys())) diff --git a/fastvideo/tests/ssim/test_flux2_similarity.py b/fastvideo/tests/ssim/test_flux2_similarity.py index 581f47e71..880fadf83 100644 --- a/fastvideo/tests/ssim/test_flux2_similarity.py +++ b/fastvideo/tests/ssim/test_flux2_similarity.py @@ -24,11 +24,9 @@ import torch from fastvideo.logger import init_logger from fastvideo.tests.ssim.inference_similarity_utils import ( - resolve_inference_device_reference_folder, -) + resolve_inference_device_reference_folder, ) from fastvideo.tests.ssim.latent_similarity_utils import ( - run_text_to_latent_similarity_test, -) + run_text_to_latent_similarity_test, ) logger = init_logger(__name__) @@ -107,8 +105,6 @@ def test_flux2_similarity( "vae_cpu_offload": True, "text_encoder_cpu_offload": True, "pin_cpu_memory": False, - "override_pipeline_cls_name": ( - "Flux2KleinPipeline" if "klein" in model_id.lower() else "Flux2Pipeline" - ), + "override_pipeline_cls_name": ("Flux2KleinPipeline" if "klein" in model_id.lower() else "Flux2Pipeline"), }, ) diff --git a/fastvideo/tests/ssim/test_flux_t2i_similarity.py b/fastvideo/tests/ssim/test_flux_t2i_similarity.py index dad9984cc..5d78feab4 100644 --- a/fastvideo/tests/ssim/test_flux_t2i_similarity.py +++ b/fastvideo/tests/ssim/test_flux_t2i_similarity.py @@ -10,8 +10,7 @@ import torch from fastvideo.api.flux import FluxSamplingParam from fastvideo.logger import init_logger from fastvideo.tests.ssim.inference_similarity_utils import ( - run_text_to_video_similarity_test, -) + run_text_to_video_similarity_test, ) from fastvideo.tests.ssim.reference_utils import ( get_cuda_device_name, resolve_device_reference_folder, @@ -103,10 +102,8 @@ def test_flux_t2i_similarity( ) -> None: is_hf_repo = "/" in FLUX_MODEL_PATH and not FLUX_MODEL_PATH.startswith("/") if not is_hf_repo and not os.path.isdir(FLUX_MODEL_PATH): - pytest.skip( - f"FLUX weights not found at {FLUX_MODEL_PATH} " - f"(set FLUX_T2I_MODEL_DIR to override)" - ) + pytest.skip(f"FLUX weights not found at {FLUX_MODEL_PATH} " + f"(set FLUX_T2I_MODEL_DIR to override)") run_text_to_video_similarity_test( logger=logger, diff --git a/fastvideo/tests/ssim/test_gamecraft_similarity.py b/fastvideo/tests/ssim/test_gamecraft_similarity.py index 42269d42a..a11ff4b2c 100644 --- a/fastvideo/tests/ssim/test_gamecraft_similarity.py +++ b/fastvideo/tests/ssim/test_gamecraft_similarity.py @@ -90,8 +90,7 @@ GAMECRAFT_T2V_PARAMS = { "negative_prompt": "", } -_GAMECRAFT_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - _GAMECRAFT_MODEL_PATH) +_GAMECRAFT_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(_GAMECRAFT_MODEL_PATH) GAMECRAFT_T2V_FULL_QUALITY_PARAMS = { "num_gpus": GAMECRAFT_T2V_PARAMS["num_gpus"], "model_path": GAMECRAFT_T2V_PARAMS["model_path"], @@ -108,10 +107,8 @@ GAMECRAFT_T2V_FULL_QUALITY_PARAMS = { GAMECRAFT_I2V_PARAMS = { **GAMECRAFT_T2V_PARAMS, - "image_path": ( - "https://huggingface.co/datasets/huggingface/documentation-images/" - "resolve/main/diffusers/astronaut.jpg" - ), + "image_path": ("https://huggingface.co/datasets/huggingface/documentation-images/" + "resolve/main/diffusers/astronaut.jpg"), } GAMECRAFT_I2V_FULL_QUALITY_PARAMS = { **GAMECRAFT_T2V_FULL_QUALITY_PARAMS, @@ -205,16 +202,12 @@ def test_gamecraft_t2v_similarity(prompt, ATTENTION_BACKEND, model_id): generator: VideoGenerator | None = None try: - generator = VideoGenerator.from_pretrained( - model_path=BASE_PARAMS["model_path"], **init_kwargs - ) + generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs) generator.generate_video(prompt, **generation_kwargs) finally: _shutdown_executor(generator) - assert os.path.exists(output_dir), ( - f"Output video was not generated at {output_dir}" - ) + assert os.path.exists(output_dir), (f"Output video was not generated at {output_dir}") # Compare to reference reference_folder = build_reference_folder_path( @@ -226,9 +219,7 @@ def test_gamecraft_t2v_similarity(prompt, ATTENTION_BACKEND, model_id): if not os.path.exists(reference_folder): logger.error("Reference folder missing") - raise FileNotFoundError( - f"Reference video folder does not exist: {reference_folder}" - ) + raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") reference_video_name = None for filename in os.listdir(reference_folder): @@ -237,21 +228,15 @@ def test_gamecraft_t2v_similarity(prompt, ATTENTION_BACKEND, model_id): break if not reference_video_name: - logger.error( - f"Reference video not found for prompt: {prompt} " - f"with backend: {ATTENTION_BACKEND}" - ) + logger.error(f"Reference video not found for prompt: {prompt} " + f"with backend: {ATTENTION_BACKEND}") raise FileNotFoundError("Reference video missing") reference_video_path = os.path.join(reference_folder, reference_video_name) generated_video_path = os.path.join(output_dir, output_video_name) - logger.info( - f"Computing SSIM between {reference_video_path} and {generated_video_path}" - ) - ssim_values = compute_video_ssim_torchvision( - reference_video_path, generated_video_path, use_ms_ssim=True - ) + logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") + ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) mean_ssim = ssim_values[0] logger.info(f"SSIM mean value: {mean_ssim}") @@ -270,10 +255,8 @@ def test_gamecraft_t2v_similarity(prompt, ATTENTION_BACKEND, model_id): logger.error("Failed to write SSIM results to file") min_acceptable_ssim = 0.93 - assert mean_ssim >= min_acceptable_ssim, ( - f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}" - ) + assert mean_ssim >= min_acceptable_ssim, (f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " + f"for {model_id} with backend {ATTENTION_BACKEND}") # --------------------------------------------------------------------------- @@ -342,16 +325,12 @@ def test_gamecraft_i2v_similarity(prompt, ATTENTION_BACKEND, model_id): generator: VideoGenerator | None = None try: - generator = VideoGenerator.from_pretrained( - model_path=BASE_PARAMS["model_path"], **init_kwargs - ) + generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs) generator.generate_video(prompt, **generation_kwargs) finally: _shutdown_executor(generator) - assert os.path.exists(output_dir), ( - f"Output video was not generated at {output_dir}" - ) + assert os.path.exists(output_dir), (f"Output video was not generated at {output_dir}") # Compare to reference reference_folder = build_reference_folder_path( @@ -363,9 +342,7 @@ def test_gamecraft_i2v_similarity(prompt, ATTENTION_BACKEND, model_id): if not os.path.exists(reference_folder): logger.error("Reference folder missing") - raise FileNotFoundError( - f"Reference video folder does not exist: {reference_folder}" - ) + raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") reference_video_name = None for filename in os.listdir(reference_folder): @@ -374,21 +351,15 @@ def test_gamecraft_i2v_similarity(prompt, ATTENTION_BACKEND, model_id): break if not reference_video_name: - logger.error( - f"Reference video not found for prompt: {prompt} " - f"with backend: {ATTENTION_BACKEND}" - ) + logger.error(f"Reference video not found for prompt: {prompt} " + f"with backend: {ATTENTION_BACKEND}") raise FileNotFoundError("Reference video missing") reference_video_path = os.path.join(reference_folder, reference_video_name) generated_video_path = os.path.join(output_dir, output_video_name) - logger.info( - f"Computing SSIM between {reference_video_path} and {generated_video_path}" - ) - ssim_values = compute_video_ssim_torchvision( - reference_video_path, generated_video_path, use_ms_ssim=True - ) + logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") + ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) mean_ssim = ssim_values[0] logger.info(f"SSIM mean value: {mean_ssim}") @@ -407,7 +378,5 @@ def test_gamecraft_i2v_similarity(prompt, ATTENTION_BACKEND, model_id): logger.error("Failed to write SSIM results to file") min_acceptable_ssim = 0.93 - assert mean_ssim >= min_acceptable_ssim, ( - f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}" - ) + assert mean_ssim >= min_acceptable_ssim, (f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " + f"for {model_id} with backend {ATTENTION_BACKEND}") diff --git a/fastvideo/tests/ssim/test_gen3c_similarity.py b/fastvideo/tests/ssim/test_gen3c_similarity.py index 6895aaa2e..07f41761c 100644 --- a/fastvideo/tests/ssim/test_gen3c_similarity.py +++ b/fastvideo/tests/ssim/test_gen3c_similarity.py @@ -71,8 +71,7 @@ else: GEN3C_T2V_PARAMS = { "num_gpus": 1, - "model_path": os.getenv("GEN3C_MODEL_PATH", - "FastVideo/GEN3C-Cosmos-7B-Diffusers"), + "model_path": os.getenv("GEN3C_MODEL_PATH", "FastVideo/GEN3C-Cosmos-7B-Diffusers"), "height": 720, "width": 1280, "num_frames": 121, @@ -99,7 +98,6 @@ TEST_PROMPTS = [ BASELINE_VIDEO_NAME = "gen3c_ssim_baseline.mp4" CANDIDATE_VIDEO_NAME = "gen3c_ssim_candidate.mp4" - # --------------------------------------------------------------------------- # Test # --------------------------------------------------------------------------- @@ -130,24 +128,18 @@ def test_gen3c_inference_similarity(prompt, ATTENTION_BACKEND, model_id): # Guard common misconfigurations to keep CI behavior explicit. if model_path.lower() == "nvidia/gen3c-cosmos-7b": - pytest.skip( - "nvidia/GEN3C-Cosmos-7B is the official raw checkpoint repo, not Diffusers format. " - "Use GEN3C_MODEL_PATH=FastVideo/GEN3C-Cosmos-7B-Diffusers or a local converted path." - ) + pytest.skip("nvidia/GEN3C-Cosmos-7B is the official raw checkpoint repo, not Diffusers format. " + "Use GEN3C_MODEL_PATH=FastVideo/GEN3C-Cosmos-7B-Diffusers or a local converted path.") local_like = model_path.startswith(("/", "./", "../")) if local_like and not os.path.exists(model_path): - pytest.skip( - f"Local GEN3C model path not found: {model_path}. " - "Set GEN3C_MODEL_PATH to a valid local path or HF Diffusers repo id." - ) + pytest.skip(f"Local GEN3C model path not found: {model_path}. " + "Set GEN3C_MODEL_PATH to a valid local path or HF Diffusers repo id.") if os.path.exists(model_path): model_index_path = os.path.join(model_path, "model_index.json") if not os.path.exists(model_index_path): - pytest.skip( - f"GEN3C_MODEL_PATH is not Diffusers-format (missing model_index.json): {model_path}" - ) + pytest.skip(f"GEN3C_MODEL_PATH is not Diffusers-format (missing model_index.json): {model_path}") init_kwargs = { "num_gpus": BASE_PARAMS["num_gpus"], @@ -171,10 +163,8 @@ def test_gen3c_inference_similarity(prompt, ATTENTION_BACKEND, model_id): } if not os.path.exists(generation_kwargs["image_path"]): - pytest.skip( - f"GEN3C test image not found: {generation_kwargs['image_path']}. " - "Set GEN3C_TEST_IMAGE_PATH to a valid local image." - ) + pytest.skip(f"GEN3C test image not found: {generation_kwargs['image_path']}. " + "Set GEN3C_TEST_IMAGE_PATH to a valid local image.") # Keep local reruns deterministic: remove prior candidate outputs so # VideoGenerator does not auto-suffix (_1, _2, ...). @@ -182,9 +172,7 @@ def test_gen3c_inference_similarity(prompt, ATTENTION_BACKEND, model_id): for stale_video in glob.glob(stale_pattern): os.remove(stale_video) - generator = VideoGenerator.from_pretrained( - model_path=model_path, **init_kwargs - ) + generator = VideoGenerator.from_pretrained(model_path=model_path, **init_kwargs) generator.generate_video(prompt, **generation_kwargs) if isinstance(generator.executor, MultiprocExecutor): @@ -192,26 +180,18 @@ def test_gen3c_inference_similarity(prompt, ATTENTION_BACKEND, model_id): assert os.path.exists(output_dir), f"Output not generated at {output_dir}" - reference_folder = os.path.join( - script_dir, device_reference_folder, model_id, ATTENTION_BACKEND - ) + reference_folder = os.path.join(script_dir, device_reference_folder, model_id, ATTENTION_BACKEND) if not os.path.exists(reference_folder): - raise FileNotFoundError( - f"Reference video folder does not exist: {reference_folder}" - ) + raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") reference_video_path = os.path.join(reference_folder, BASELINE_VIDEO_NAME) if not os.path.exists(reference_video_path): - raise FileNotFoundError( - f"Reference video not found: {reference_video_path}" - ) + raise FileNotFoundError(f"Reference video not found: {reference_video_path}") generated_video_path = os.path.join(output_dir, output_video_name) logger.info(f"Computing SSIM: {reference_video_path} vs {generated_video_path}") - ssim_values = compute_video_ssim_torchvision( - reference_video_path, generated_video_path, use_ms_ssim=True - ) + ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) mean_ssim = ssim_values[0] logger.info(f"GEN3C SSIM mean: {mean_ssim}") @@ -228,5 +208,4 @@ def test_gen3c_inference_similarity(prompt, ATTENTION_BACKEND, model_id): # GEN3C SSIM threshold for stable L40S reference comparisons. min_acceptable_ssim = 0.93 assert mean_ssim >= min_acceptable_ssim, ( - f"SSIM {mean_ssim:.4f} < {min_acceptable_ssim} for {model_id} / {ATTENTION_BACKEND}" - ) + f"SSIM {mean_ssim:.4f} < {min_acceptable_ssim} for {model_id} / {ATTENTION_BACKEND}") diff --git a/fastvideo/tests/ssim/test_lingbot_similarity.py b/fastvideo/tests/ssim/test_lingbot_similarity.py index 353cc68e4..a0163d839 100644 --- a/fastvideo/tests/ssim/test_lingbot_similarity.py +++ b/fastvideo/tests/ssim/test_lingbot_similarity.py @@ -35,15 +35,12 @@ def _find_lingbotworld_examples_root() -> str | None: script_dir = os.path.dirname(os.path.abspath(__file__)) repo_root = os.path.abspath(os.path.join(script_dir, "..", "..", "..")) candidates = [ - os.path.join(repo_root, "examples", "inference", "basic", - "lingbotworld_examples"), - os.path.join(repo_root, "..", "FastVideo", "examples", "inference", - "basic", "lingbotworld_examples"), + os.path.join(repo_root, "examples", "inference", "basic", "lingbotworld_examples"), + os.path.join(repo_root, "..", "FastVideo", "examples", "inference", "basic", "lingbotworld_examples"), ] for candidate in candidates: if (os.path.exists(os.path.join(candidate, "00", "poses.npy")) - and os.path.exists( - os.path.join(candidate, "00", "intrinsics.npy"))): + and os.path.exists(os.path.join(candidate, "00", "intrinsics.npy"))): return os.path.abspath(candidate) return None @@ -63,38 +60,47 @@ device_reference_folder = resolve_device_reference_folder( logger=logger, ) - LINGBOT_PARAMS = { - "model_path": "FastVideo/LingBot-World-Base-Cam-Diffusers", - "num_gpus": 2, - "height": 256, - "width": 448, - "num_frames": 45, # must be 4k+1 - "num_inference_steps": 4, - "guidance_scale": 5.0, - "guidance_scale_2": 5.0, - "embedded_cfg_scale": 6, - "flow_shift": 10.0, - "boundary_ratio": 0.947, - "seed": 42, - "fps": 16, - "spatial_scale": 8, - "example_case": "00", - "image_path": ( - "https://raw.githubusercontent.com/Robbyant/lingbot-world/main/" - "examples/00/image.jpg" - ), - "negative_prompt": ( - "画面突变，色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，" - "最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，" - "畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走，" - "镜头晃动，画面闪烁，模糊，噪点，水印，签名，文字，变形，扭曲，液化，不合逻辑的结构，卡顿，" - "PPT幻灯片感，过暗，欠曝，低对比度，霓虹灯光感，过度锐化，3D渲染感，人物，行人，游客，身体，" - "皮肤，肢体，面部特征，汽车，电线" - ), + "model_path": + "FastVideo/LingBot-World-Base-Cam-Diffusers", + "num_gpus": + 2, + "height": + 256, + "width": + 448, + "num_frames": + 45, # must be 4k+1 + "num_inference_steps": + 4, + "guidance_scale": + 5.0, + "guidance_scale_2": + 5.0, + "embedded_cfg_scale": + 6, + "flow_shift": + 10.0, + "boundary_ratio": + 0.947, + "seed": + 42, + "fps": + 16, + "spatial_scale": + 8, + "example_case": + "00", + "image_path": ("https://raw.githubusercontent.com/Robbyant/lingbot-world/main/" + "examples/00/image.jpg"), + "negative_prompt": ("画面突变，色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，" + "最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，" + "畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走，" + "镜头晃动，画面闪烁，模糊，噪点，水印，签名，文字，变形，扭曲，液化，不合逻辑的结构，卡顿，" + "PPT幻灯片感，过暗，欠曝，低对比度，霓虹灯光感，过度锐化，3D渲染感，人物，行人，游客，身体，" + "皮肤，肢体，面部特征，汽车，电线"), } -_LINGBOT_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - LINGBOT_PARAMS["model_path"]) +_LINGBOT_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(LINGBOT_PARAMS["model_path"]) LINGBOT_FULL_QUALITY_PARAMS = { "model_path": LINGBOT_PARAMS["model_path"], "num_gpus": LINGBOT_PARAMS["num_gpus"], @@ -134,15 +140,12 @@ def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): if device_reference_folder is None: pytest.skip(f"Unsupported device for LingBot SSIM test: {device_name}") if torch.cuda.device_count() < params["num_gpus"]: - pytest.skip( - f"LingBot SSIM test requires {params['num_gpus']} GPUs, " - f"but only {torch.cuda.device_count()} detected." - ) + pytest.skip(f"LingBot SSIM test requires {params['num_gpus']} GPUs, " + f"but only {torch.cuda.device_count()} detected.") examples_root = _find_lingbotworld_examples_root() if examples_root is None: - pytest.skip( - "lingbotworld_examples not found under examples/inference/basic.") + pytest.skip("lingbotworld_examples not found under examples/inference/basic.") action_path = os.path.join(examples_root, params["example_case"]) if not (os.path.exists(os.path.join(action_path, "poses.npy")) @@ -197,16 +200,14 @@ def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): generator: VideoGenerator | None = None try: - generator = VideoGenerator.from_pretrained( - model_path=params["model_path"], **init_kwargs) + generator = VideoGenerator.from_pretrained(model_path=params["model_path"], **init_kwargs) generator.generate_video(prompt, **generation_kwargs) finally: if generator is not None: generator.shutdown() generated_video_path = os.path.join(output_dir, output_video_name) - assert os.path.exists(generated_video_path), ( - f"Output video was not generated at {generated_video_path}") + assert os.path.exists(generated_video_path), (f"Output video was not generated at {generated_video_path}") reference_folder = build_reference_folder_path( script_dir, @@ -215,8 +216,7 @@ def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): ATTENTION_BACKEND, ) if not os.path.exists(reference_folder): - raise FileNotFoundError( - f"Reference video folder does not exist: {reference_folder}") + raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") reference_video_name = None for filename in os.listdir(reference_folder): @@ -224,24 +224,17 @@ def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): reference_video_name = filename break if not reference_video_name: - raise FileNotFoundError( - f"Reference video missing for prompt/backend under {reference_folder}" - ) + raise FileNotFoundError(f"Reference video missing for prompt/backend under {reference_folder}") reference_video_path = os.path.join(reference_folder, reference_video_name) - logger.info("Computing SSIM between %s and %s", reference_video_path, - generated_video_path) - ssim_values = compute_video_ssim_torchvision(reference_video_path, - generated_video_path, - use_ms_ssim=True) + logger.info("Computing SSIM between %s and %s", reference_video_path, generated_video_path) + ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) mean_ssim = ssim_values[0] logger.info("SSIM mean value: %s", mean_ssim) - write_ssim_results(output_dir, ssim_values, reference_video_path, - generated_video_path, + write_ssim_results(output_dir, ssim_values, reference_video_path, generated_video_path, params["num_inference_steps"], prompt) min_acceptable_ssim = 0.70 - assert mean_ssim >= min_acceptable_ssim, ( - f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}") + assert mean_ssim >= min_acceptable_ssim, (f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " + f"for {model_id} with backend {ATTENTION_BACKEND}") diff --git a/fastvideo/tests/ssim/test_longcat_similarity.py b/fastvideo/tests/ssim/test_longcat_similarity.py index 890789bff..aa9dd2342 100644 --- a/fastvideo/tests/ssim/test_longcat_similarity.py +++ b/fastvideo/tests/ssim/test_longcat_similarity.py @@ -49,13 +49,11 @@ device_reference_folder = resolve_device_reference_folder( ) # Common negative prompt from example scripts -NEGATIVE_PROMPT = ( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, " - "paintings, images, static, overall gray, worst quality, low quality, JPEG compression " - "residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, " - "deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, " - "three legs, many people in the background, walking backwards" -) +NEGATIVE_PROMPT = ("Bright tones, overexposed, static, blurred details, subtitles, style, works, " + "paintings, images, static, overall gray, worst quality, low quality, JPEG compression " + "residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, " + "deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, " + "three legs, many people in the background, walking backwards") # ============================================================================= # LongCat T2V Parameters (from basic_longcat_t2v.py) @@ -72,9 +70,7 @@ LONGCAT_T2V_PARAMS = { "seed": 42, "negative_prompt": NEGATIVE_PROMPT, } -_LONGCAT_T2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - LONGCAT_T2V_PARAMS["model_path"] -) +_LONGCAT_T2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(LONGCAT_T2V_PARAMS["model_path"]) LONGCAT_T2V_FULL_QUALITY_PARAMS = { "num_gpus": LONGCAT_T2V_PARAMS["num_gpus"], "model_path": LONGCAT_T2V_PARAMS["model_path"], @@ -103,9 +99,7 @@ LONGCAT_I2V_PARAMS = { "seed": 42, "negative_prompt": NEGATIVE_PROMPT, } -_LONGCAT_I2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - LONGCAT_I2V_PARAMS["model_path"] -) +_LONGCAT_I2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(LONGCAT_I2V_PARAMS["model_path"]) LONGCAT_I2V_FULL_QUALITY_PARAMS = { "num_gpus": LONGCAT_I2V_PARAMS["num_gpus"], "model_path": LONGCAT_I2V_PARAMS["model_path"], @@ -135,9 +129,7 @@ LONGCAT_VC_PARAMS = { "num_cond_frames": 13, "negative_prompt": NEGATIVE_PROMPT, } -_LONGCAT_VC_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - LONGCAT_VC_PARAMS["model_path"] -) +_LONGCAT_VC_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(LONGCAT_VC_PARAMS["model_path"]) LONGCAT_VC_FULL_QUALITY_PARAMS = { "num_gpus": LONGCAT_VC_PARAMS["num_gpus"], "model_path": LONGCAT_VC_PARAMS["model_path"], @@ -211,7 +203,7 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str): script_dir = os.path.dirname(os.path.abspath(__file__)) model_id = "LongCat-Video-T2V" - + output_dir = build_generated_output_dir( script_dir, device_reference_folder, @@ -242,16 +234,12 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str): "negative_prompt": params["negative_prompt"], } - generator = VideoGenerator.from_pretrained( - model_path=params["model_path"], **init_kwargs - ) + generator = VideoGenerator.from_pretrained(model_path=params["model_path"], **init_kwargs) generator.generate_video(prompt, **generation_kwargs) generator.shutdown() generated_video_path = os.path.join(output_dir, output_video_name) - assert os.path.exists(generated_video_path), ( - f"Output video was not generated at {generated_video_path}" - ) + assert os.path.exists(generated_video_path), (f"Output video was not generated at {generated_video_path}") # Find reference video reference_folder = build_reference_folder_path( @@ -261,9 +249,7 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str): ATTENTION_BACKEND, ) if not os.path.exists(reference_folder): - raise FileNotFoundError( - f"Reference video folder does not exist: {reference_folder}" - ) + raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") reference_video_name = None for filename in os.listdir(reference_folder): @@ -273,29 +259,22 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str): if not reference_video_name: raise FileNotFoundError( - f"Reference video not found for prompt: {prompt[:50]}... with backend: {ATTENTION_BACKEND}" - ) + f"Reference video not found for prompt: {prompt[:50]}... with backend: {ATTENTION_BACKEND}") reference_video_path = os.path.join(reference_folder, reference_video_name) logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") - ssim_values = compute_video_ssim_torchvision( - reference_video_path, generated_video_path, use_ms_ssim=True - ) + ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) mean_ssim = ssim_values[0] logger.info(f"SSIM mean value: {mean_ssim}") - write_ssim_results( - output_dir, ssim_values, reference_video_path, generated_video_path, - params["num_inference_steps"], prompt - ) + write_ssim_results(output_dir, ssim_values, reference_video_path, generated_video_path, + params["num_inference_steps"], prompt) min_acceptable_ssim = 0.90 - assert mean_ssim >= min_acceptable_ssim, ( - f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}" - ) + assert mean_ssim >= min_acceptable_ssim, (f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " + f"for {model_id} with backend {ATTENTION_BACKEND}") @pytest.mark.parametrize("prompt", I2V_TEST_PROMPTS) @@ -312,7 +291,7 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): script_dir = os.path.dirname(os.path.abspath(__file__)) model_id = "LongCat-Video-I2V" - + output_dir = build_generated_output_dir( script_dir, device_reference_folder, @@ -348,16 +327,12 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): "negative_prompt": params["negative_prompt"], } - generator = VideoGenerator.from_pretrained( - model_path=params["model_path"], **init_kwargs - ) + generator = VideoGenerator.from_pretrained(model_path=params["model_path"], **init_kwargs) generator.generate_video(prompt, **generation_kwargs) generator.shutdown() generated_video_path = os.path.join(output_dir, output_video_name) - assert os.path.exists(generated_video_path), ( - f"Output video was not generated at {generated_video_path}" - ) + assert os.path.exists(generated_video_path), (f"Output video was not generated at {generated_video_path}") # Find reference video reference_folder = build_reference_folder_path( @@ -367,9 +342,7 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): ATTENTION_BACKEND, ) if not os.path.exists(reference_folder): - raise FileNotFoundError( - f"Reference video folder does not exist: {reference_folder}" - ) + raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") reference_video_name = None for filename in os.listdir(reference_folder): @@ -379,29 +352,22 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): if not reference_video_name: raise FileNotFoundError( - f"Reference video not found for prompt: {prompt[:50]}... with backend: {ATTENTION_BACKEND}" - ) + f"Reference video not found for prompt: {prompt[:50]}... with backend: {ATTENTION_BACKEND}") reference_video_path = os.path.join(reference_folder, reference_video_name) logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") - ssim_values = compute_video_ssim_torchvision( - reference_video_path, generated_video_path, use_ms_ssim=True - ) + ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) mean_ssim = ssim_values[0] logger.info(f"SSIM mean value: {mean_ssim}") - write_ssim_results( - output_dir, ssim_values, reference_video_path, generated_video_path, - params["num_inference_steps"], prompt - ) + write_ssim_results(output_dir, ssim_values, reference_video_path, generated_video_path, + params["num_inference_steps"], prompt) min_acceptable_ssim = 0.90 - assert mean_ssim >= min_acceptable_ssim, ( - f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}" - ) + assert mean_ssim >= min_acceptable_ssim, (f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " + f"for {model_id} with backend {ATTENTION_BACKEND}") @pytest.mark.parametrize("prompt", VC_TEST_PROMPTS) @@ -418,7 +384,7 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str): script_dir = os.path.dirname(os.path.abspath(__file__)) model_id = "LongCat-Video-VC" - + output_dir = build_generated_output_dir( script_dir, device_reference_folder, @@ -431,7 +397,7 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str): # Get video path for this prompt prompt_idx = VC_TEST_PROMPTS.index(prompt) video_path = _resolve_asset_path(VC_VIDEO_PATHS[prompt_idx]) - + if not os.path.exists(video_path): pytest.skip(f"Input video not found at {video_path}") @@ -459,16 +425,12 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str): "negative_prompt": params["negative_prompt"], } - generator = VideoGenerator.from_pretrained( - model_path=params["model_path"], **init_kwargs - ) + generator = VideoGenerator.from_pretrained(model_path=params["model_path"], **init_kwargs) generator.generate_video(prompt, **generation_kwargs) generator.shutdown() generated_video_path = os.path.join(output_dir, output_video_name) - assert os.path.exists(generated_video_path), ( - f"Output video was not generated at {generated_video_path}" - ) + assert os.path.exists(generated_video_path), (f"Output video was not generated at {generated_video_path}") # Find reference video reference_folder = build_reference_folder_path( @@ -478,9 +440,7 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str): ATTENTION_BACKEND, ) if not os.path.exists(reference_folder): - raise FileNotFoundError( - f"Reference video folder does not exist: {reference_folder}" - ) + raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") reference_video_name = None for filename in os.listdir(reference_folder): @@ -490,26 +450,19 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str): if not reference_video_name: raise FileNotFoundError( - f"Reference video not found for prompt: {prompt[:50]}... with backend: {ATTENTION_BACKEND}" - ) + f"Reference video not found for prompt: {prompt[:50]}... with backend: {ATTENTION_BACKEND}") reference_video_path = os.path.join(reference_folder, reference_video_name) logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") - ssim_values = compute_video_ssim_torchvision( - reference_video_path, generated_video_path, use_ms_ssim=True - ) + ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) mean_ssim = ssim_values[0] logger.info(f"SSIM mean value: {mean_ssim}") - write_ssim_results( - output_dir, ssim_values, reference_video_path, generated_video_path, - params["num_inference_steps"], prompt - ) + write_ssim_results(output_dir, ssim_values, reference_video_path, generated_video_path, + params["num_inference_steps"], prompt) min_acceptable_ssim = 0.90 - assert mean_ssim >= min_acceptable_ssim, ( - f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}" - ) + assert mean_ssim >= min_acceptable_ssim, (f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " + f"for {model_id} with backend {ATTENTION_BACKEND}") diff --git a/fastvideo/tests/ssim/test_ltx2_similarity.py b/fastvideo/tests/ssim/test_ltx2_similarity.py index 41a7c9d43..73031a242 100644 --- a/fastvideo/tests/ssim/test_ltx2_similarity.py +++ b/fastvideo/tests/ssim/test_ltx2_similarity.py @@ -27,11 +27,9 @@ import pytest from fastvideo.api.sampling_param import SamplingParam from fastvideo.logger import init_logger from fastvideo.tests.ssim.inference_similarity_utils import ( - resolve_inference_device_reference_folder, -) + resolve_inference_device_reference_folder, ) from fastvideo.tests.ssim.latent_similarity_utils import ( - run_text_to_latent_similarity_test, -) + run_text_to_latent_similarity_test, ) logger = init_logger(__name__) @@ -62,16 +60,14 @@ LTX2_DISTILLED_REFERENCE_GUIDANCE_OVERRIDES = { "ltx2_stg_scale_video": 1.0, "ltx2_stg_scale_audio": 1.0, } -_LTX2_DISTILLED_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - LTX2_DISTILLED_PARAMS["model_path"]) +_LTX2_DISTILLED_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(LTX2_DISTILLED_PARAMS["model_path"]) LTX2_DISTILLED_FULL_QUALITY_PARAMS = { "num_gpus": LTX2_DISTILLED_PARAMS["num_gpus"], "model_path": LTX2_DISTILLED_PARAMS["model_path"], "height": _LTX2_DISTILLED_FULL_QUALITY_DEFAULTS.height, "width": _LTX2_DISTILLED_FULL_QUALITY_DEFAULTS.width, "num_frames": _LTX2_DISTILLED_FULL_QUALITY_DEFAULTS.num_frames, - "num_inference_steps": - _LTX2_DISTILLED_FULL_QUALITY_DEFAULTS.num_inference_steps, + "num_inference_steps": _LTX2_DISTILLED_FULL_QUALITY_DEFAULTS.num_inference_steps, "guidance_scale": _LTX2_DISTILLED_FULL_QUALITY_DEFAULTS.guidance_scale, "seed": _LTX2_DISTILLED_FULL_QUALITY_DEFAULTS.seed, "sp_size": LTX2_DISTILLED_PARAMS["sp_size"], diff --git a/fastvideo/tests/ssim/test_matrixgame2_similarity.py b/fastvideo/tests/ssim/test_matrixgame2_similarity.py index 34a94be3a..b7c21bdd6 100644 --- a/fastvideo/tests/ssim/test_matrixgame2_similarity.py +++ b/fastvideo/tests/ssim/test_matrixgame2_similarity.py @@ -48,8 +48,7 @@ MATRIXGAME_PARAMS = { "seed": 1024, "keyboard_dim": 4, } -_MATRIXGAME_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - MATRIXGAME_PARAMS["model_path"]) +_MATRIXGAME_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(MATRIXGAME_PARAMS["model_path"]) MATRIXGAME_FULL_QUALITY_PARAMS = { "num_gpus": MATRIXGAME_PARAMS["num_gpus"], "model_path": MATRIXGAME_PARAMS["model_path"], @@ -62,7 +61,6 @@ MATRIXGAME_FULL_QUALITY_PARAMS = { "keyboard_dim": MATRIXGAME_PARAMS["keyboard_dim"], } - MODEL_TO_PARAMS = { "Matrix-Game-2.0-Diffusers-Base": MATRIXGAME_PARAMS, } @@ -111,9 +109,9 @@ def test_matrixgame2_similarity(prompt, ATTENTION_BACKEND, model_id): num_inference_steps = BASE_PARAMS["num_inference_steps"] # Create action conditions for Matrix-Game 2.0 - actions = create_action_presets( - BASE_PARAMS["num_frames"], keyboard_dim=BASE_PARAMS["keyboard_dim"], seed=BASE_PARAMS["seed"] - ) + actions = create_action_presets(BASE_PARAMS["num_frames"], + keyboard_dim=BASE_PARAMS["keyboard_dim"], + seed=BASE_PARAMS["seed"]) latent_frames = (BASE_PARAMS["num_frames"] - 1) // 4 + 1 grid_sizes = torch.tensor([latent_frames, 44, 80]) diff --git a/fastvideo/tests/ssim/test_matrixgame3_similarity.py b/fastvideo/tests/ssim/test_matrixgame3_similarity.py index ff3753f41..b383698b1 100644 --- a/fastvideo/tests/ssim/test_matrixgame3_similarity.py +++ b/fastvideo/tests/ssim/test_matrixgame3_similarity.py @@ -35,7 +35,6 @@ device_reference_folder = resolve_device_reference_folder( logger=logger, ) - MATRIXGAME3_PARAMS = { "num_gpus": 1, "model_path": "FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers", @@ -46,8 +45,7 @@ MATRIXGAME3_PARAMS = { "guidance_scale": 5.0, "seed": 42, } -_MATRIXGAME3_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - MATRIXGAME3_PARAMS["model_path"]) +_MATRIXGAME3_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(MATRIXGAME3_PARAMS["model_path"]) MATRIXGAME3_FULL_QUALITY_PARAMS = { "num_gpus": MATRIXGAME3_PARAMS["num_gpus"], "model_path": MATRIXGAME3_PARAMS["model_path"], @@ -59,7 +57,6 @@ MATRIXGAME3_FULL_QUALITY_PARAMS = { "seed": _MATRIXGAME3_FULL_QUALITY_DEFAULTS.seed, } - MODEL_TO_PARAMS = { "Matrix-Game-3.0-Diffusers-Base": MATRIXGAME3_PARAMS, } diff --git a/fastvideo/tests/ssim/test_reference_videos_cli.py b/fastvideo/tests/ssim/test_reference_videos_cli.py index 20a51d4e6..ec6ab8417 100644 --- a/fastvideo/tests/ssim/test_reference_videos_cli.py +++ b/fastvideo/tests/ssim/test_reference_videos_cli.py @@ -73,25 +73,11 @@ def _fake_hf(monkeypatch): def test_upload_draft_reference_artifact_uses_drafts_prefix(tmp_path): ssim_dir = tmp_path / "ssim" - generated = ( - ssim_dir - / "generated_videos" - / "default" - / "L40S_reference_videos" - / "model-a" - / "flash" - / "sample.mp4" - ) + generated = (ssim_dir / "generated_videos" / "default" / "L40S_reference_videos" / "model-a" / "flash" / + "sample.mp4") generated.parent.mkdir(parents=True) generated.write_text("video", encoding="utf-8") - reference_folder = ( - ssim_dir - / "reference_videos" - / "default" - / "L40S_reference_videos" - / "model-a" - / "flash" - ) + reference_folder = (ssim_dir / "reference_videos" / "default" / "L40S_reference_videos" / "model-a" / "flash") draft_path = reference_videos_cli.upload_draft_reference_artifact( repo_id="FastVideo/ssim-reference-videos", diff --git a/fastvideo/tests/ssim/test_sd35_similarity.py b/fastvideo/tests/ssim/test_sd35_similarity.py index ffe0ab1c2..ba64fa029 100644 --- a/fastvideo/tests/ssim/test_sd35_similarity.py +++ b/fastvideo/tests/ssim/test_sd35_similarity.py @@ -9,8 +9,7 @@ import torch from fastvideo.api.sampling_param import SamplingParam from fastvideo.logger import init_logger from fastvideo.tests.ssim.inference_similarity_utils import ( - run_text_to_video_similarity_test, -) + run_text_to_video_similarity_test, ) from fastvideo.tests.ssim.reference_utils import ( get_cuda_device_name, resolve_device_reference_folder, @@ -61,8 +60,7 @@ SD35_PARAMS = { "neg_prompt": "", } -_SD35_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - SD35_MODEL_PATH) +_SD35_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(SD35_MODEL_PATH) SD35_FULL_QUALITY_PARAMS = { "num_gpus": 1, "model_path": SD35_MODEL_PATH, diff --git a/fastvideo/tests/ssim/test_stable_audio_similarity.py b/fastvideo/tests/ssim/test_stable_audio_similarity.py index acf9ae0f6..340997ea5 100644 --- a/fastvideo/tests/ssim/test_stable_audio_similarity.py +++ b/fastvideo/tests/ssim/test_stable_audio_similarity.py @@ -31,8 +31,7 @@ import pytest from fastvideo.api.sampling_param import SamplingParam from fastvideo.logger import init_logger from fastvideo.tests.ssim.inference_similarity_utils import ( - resolve_inference_device_reference_folder, -) + resolve_inference_device_reference_folder, ) from fastvideo.tests.ssim.latent_similarity_utils import ( AUDIO_FIRST_8_TIMESTEPS_SPEC, run_text_to_latent_similarity_test, diff --git a/fastvideo/tests/ssim/test_turbodiffusion_similarity.py b/fastvideo/tests/ssim/test_turbodiffusion_similarity.py index 2c87d4244..acd471360 100644 --- a/fastvideo/tests/ssim/test_turbodiffusion_similarity.py +++ b/fastvideo/tests/ssim/test_turbodiffusion_similarity.py @@ -41,8 +41,7 @@ TURBODIFFUSION_PARAMS = { "tp_size": 1, "fps": 24, } -_TURBODIFFUSION_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - TURBODIFFUSION_PARAMS["model_path"]) +_TURBODIFFUSION_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(TURBODIFFUSION_PARAMS["model_path"]) TURBODIFFUSION_FULL_QUALITY_PARAMS = { "num_gpus": TURBODIFFUSION_PARAMS["num_gpus"], "model_path": TURBODIFFUSION_PARAMS["model_path"], @@ -110,8 +109,7 @@ TURBODIFFUSION_I2V_PARAMS = { "tp_size": 1, "fps": 24, } -_TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - TURBODIFFUSION_I2V_PARAMS["model_path"]) +_TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(TURBODIFFUSION_I2V_PARAMS["model_path"]) TURBODIFFUSION_I2V_FULL_QUALITY_PARAMS = { "num_gpus": TURBODIFFUSION_I2V_PARAMS["num_gpus"], "model_path": TURBODIFFUSION_I2V_PARAMS["model_path"], diff --git a/fastvideo/tests/ssim/test_wan_i2v_similarity.py b/fastvideo/tests/ssim/test_wan_i2v_similarity.py index 0f24b9add..6ff07c2ab 100644 --- a/fastvideo/tests/ssim/test_wan_i2v_similarity.py +++ b/fastvideo/tests/ssim/test_wan_i2v_similarity.py @@ -17,32 +17,43 @@ REQUIRED_GPUS = 2 device_reference_folder = resolve_inference_device_reference_folder(logger) WAN_I2V_PARAMS = { - "num_gpus": 2, - "model_path": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", - "height": 480, - "width": 832, - "num_frames": 45, - "num_inference_steps": 2, - "guidance_scale": 5.0, - "embedded_cfg_scale": 6, - "flow_shift": 7.0, - "seed": 1024, - "sp_size": 2, - "tp_size": 1, - "vae_sp": True, - "fps": 24, - "neg_prompt": ( - "Bright tones, overexposed, static, blurred details, subtitles, style, " - "works, paintings, images, static, overall gray, worst quality, low " - "quality, JPEG compression residue, ugly, incomplete, extra fingers, " - "poorly drawn hands, poorly drawn faces, deformed, disfigured, " - "misshapen limbs, fused fingers, still picture, messy background, " - "three legs, many people in the background, walking backwards" - ), - "text-encoder-precision": ("fp32",), + "num_gpus": + 2, + "model_path": + "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", + "height": + 480, + "width": + 832, + "num_frames": + 45, + "num_inference_steps": + 2, + "guidance_scale": + 5.0, + "embedded_cfg_scale": + 6, + "flow_shift": + 7.0, + "seed": + 1024, + "sp_size": + 2, + "tp_size": + 1, + "vae_sp": + True, + "fps": + 24, + "neg_prompt": ("Bright tones, overexposed, static, blurred details, subtitles, style, " + "works, paintings, images, static, overall gray, worst quality, low " + "quality, JPEG compression residue, ugly, incomplete, extra fingers, " + "poorly drawn hands, poorly drawn faces, deformed, disfigured, " + "misshapen limbs, fused fingers, still picture, messy background, " + "three legs, many people in the background, walking backwards"), + "text-encoder-precision": ("fp32", ), } -_WAN_I2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers") +_WAN_I2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained("Wan-AI/Wan2.1-I2V-14B-480P-Diffusers") WAN_I2V_FULL_QUALITY_PARAMS = { "num_gpus": WAN_I2V_PARAMS["num_gpus"], "model_path": WAN_I2V_PARAMS["model_path"], diff --git a/fastvideo/tests/ssim/test_wan_t2v_similarity.py b/fastvideo/tests/ssim/test_wan_t2v_similarity.py index 51ec7baf2..699b6a765 100644 --- a/fastvideo/tests/ssim/test_wan_t2v_similarity.py +++ b/fastvideo/tests/ssim/test_wan_t2v_similarity.py @@ -17,32 +17,43 @@ REQUIRED_GPUS = 2 device_reference_folder = resolve_inference_device_reference_folder(logger) WAN_T2V_PARAMS = { - "num_gpus": 2, - "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "height": 480, - "width": 832, - "num_frames": 45, - "num_inference_steps": 4, - "guidance_scale": 3, - "embedded_cfg_scale": 6, - "flow_shift": 7.0, - "seed": 1024, - "sp_size": 2, - "tp_size": 1, - "vae_sp": True, - "fps": 24, - "neg_prompt": ( - "Bright tones, overexposed, static, blurred details, subtitles, style, " - "works, paintings, images, static, overall gray, worst quality, low " - "quality, JPEG compression residue, ugly, incomplete, extra fingers, " - "poorly drawn hands, poorly drawn faces, deformed, disfigured, " - "misshapen limbs, fused fingers, still picture, messy background, " - "three legs, many people in the background, walking backwards" - ), - "text-encoder-precision": ("fp32",), + "num_gpus": + 2, + "model_path": + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "height": + 480, + "width": + 832, + "num_frames": + 45, + "num_inference_steps": + 4, + "guidance_scale": + 3, + "embedded_cfg_scale": + 6, + "flow_shift": + 7.0, + "seed": + 1024, + "sp_size": + 2, + "tp_size": + 1, + "vae_sp": + True, + "fps": + 24, + "neg_prompt": ("Bright tones, overexposed, static, blurred details, subtitles, style, " + "works, paintings, images, static, overall gray, worst quality, low " + "quality, JPEG compression residue, ugly, incomplete, extra fingers, " + "poorly drawn hands, poorly drawn faces, deformed, disfigured, " + "misshapen limbs, fused fingers, still picture, messy background, " + "three legs, many people in the background, walking backwards"), + "text-encoder-precision": ("fp32", ), } -_WAN_T2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( - "Wan-AI/Wan2.1-T2V-1.3B-Diffusers") +_WAN_T2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") WAN_T2V_FULL_QUALITY_PARAMS = { "num_gpus": WAN_T2V_PARAMS["num_gpus"], "model_path": WAN_T2V_PARAMS["model_path"], diff --git a/fastvideo/tests/stages/test_kandinsky5_dmd_stage_backend_engages.py b/fastvideo/tests/stages/test_kandinsky5_dmd_stage_backend_engages.py index 12d8bca75..075530022 100644 --- a/fastvideo/tests/stages/test_kandinsky5_dmd_stage_backend_engages.py +++ b/fastvideo/tests/stages/test_kandinsky5_dmd_stage_backend_engages.py @@ -173,10 +173,9 @@ def test_raw_transformer_without_context_bypasses_backend(monkeypatch): ).sample assert out.shape == (1, 2, 8, 8, 4) - assert len(backend_calls) == 0, ( - "expected the missing-forward-context fallback to bypass the backend impl; " - "if this now fails, Kandinsky5Attention's fallback behavior changed and the " - "stage guard docs/tests should be revisited") + assert len(backend_calls) == 0, ("expected the missing-forward-context fallback to bypass the backend impl; " + "if this now fails, Kandinsky5Attention's fallback behavior changed and the " + "stage guard docs/tests should be revisited") finally: _cached_get_attn_backend.cache_clear() diff --git a/fastvideo/tests/stages/test_kandinsky5_target_dtype.py b/fastvideo/tests/stages/test_kandinsky5_target_dtype.py index 0c5b763be..b07135669 100644 --- a/fastvideo/tests/stages/test_kandinsky5_target_dtype.py +++ b/fastvideo/tests/stages/test_kandinsky5_target_dtype.py @@ -78,9 +78,8 @@ def test_resolve_target_dtype_honors_explicit_fp32_for_plain_transformer(): resolved = stage._resolve_target_dtype(_fastvideo_args("fp32")) - assert resolved == torch.float32, ( - "an explicit fp32 pipeline_config must not be silently cast to bf16 for a " - "plain (non-FSDP) transformer") + assert resolved == torch.float32, ("an explicit fp32 pipeline_config must not be silently cast to bf16 for a " + "plain (non-FSDP) transformer") def test_resolve_target_dtype_honors_explicit_fp16_for_plain_transformer(): @@ -110,9 +109,8 @@ def test_resolve_target_dtype_fsdp_reads_policy_not_parameter_storage(mixed_prec resolved = stage._resolve_target_dtype(_fastvideo_args("fp16")) - assert resolved == torch.bfloat16, ( - "FSDP compute dtype comes from the MixedPrecisionPolicy param_dtype, not from " - "the dtype the parameters happen to be stored in") + assert resolved == torch.bfloat16, ("FSDP compute dtype comes from the MixedPrecisionPolicy param_dtype, not from " + "the dtype the parameters happen to be stored in") def test_resolve_target_dtype_fsdp_follows_a_non_default_policy(mixed_precision_state_reset): diff --git a/fastvideo/tests/stages/test_magi_human_sr_latent_preparation.py b/fastvideo/tests/stages/test_magi_human_sr_latent_preparation.py index 3237a72c5..91e0917b8 100644 --- a/fastvideo/tests/stages/test_magi_human_sr_latent_preparation.py +++ b/fastvideo/tests/stages/test_magi_human_sr_latent_preparation.py @@ -33,8 +33,7 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch def _make_stage() -> MagiHumanSRLatentPreparationStage: """Bypass __init__: only set the fields the T2V forward() path reads.""" - stage = MagiHumanSRLatentPreparationStage.__new__( - MagiHumanSRLatentPreparationStage) + stage = MagiHumanSRLatentPreparationStage.__new__(MagiHumanSRLatentPreparationStage) # vae + video_processor are only used by `_encode_image` (TI2V path); # T2V skips that branch when `batch.image_latent is None`. stage.vae = None @@ -47,8 +46,7 @@ def _make_stage() -> MagiHumanSRLatentPreparationStage: stage.sr_audio_noise_scale = 0.7 stage.sr_height = 512 stage.sr_width = 896 - stage.sigmas = ZeroSNRDDPMDiscretization()( - 1000, do_append_zero=False, flip=True) + stage.sigmas = ZeroSNRDDPMDiscretization()(1000, do_append_zero=False, flip=True) return stage diff --git a/fastvideo/tests/stages/test_text_encoding.py b/fastvideo/tests/stages/test_text_encoding.py index 87c55f7f1..754144b61 100644 --- a/fastvideo/tests/stages/test_text_encoding.py +++ b/fastvideo/tests/stages/test_text_encoding.py @@ -9,11 +9,15 @@ from fastvideo.configs.pipelines.base import PipelineConfig from fastvideo.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig, BaseEncoderOutput from fastvideo.pipelines.stages.text_encoding import TextEncodingStage + class TensorDict(dict): + def to(self, device): return TensorDict({k: v.to(device) for k, v in self.items()}) + class FakeTokenizer: + def __init__(self): self.last_chat_template_kwargs = None @@ -31,6 +35,7 @@ class FakeTokenizer: class FakeChatTokenizer: + def __init__(self): self.last_messages = None self.last_kwargs = None @@ -48,7 +53,9 @@ class FakeChatTokenizer: "attention_mask": torch.ones(B, seq_len, dtype=torch.long), }) + class FakeTextEncoder(torch.nn.Module): + def __init__(self, hidden_size=8): super().__init__() self.hidden_size = hidden_size @@ -65,8 +72,8 @@ class FakeTextEncoder(torch.nn.Module): device=input_ids.device, ).view(B, T, self.hidden_size) hidden_states = (last_hidden_state, ) if output_hidden_states else None - return types.SimpleNamespace(last_hidden_state=last_hidden_state, - hidden_states=hidden_states) + return types.SimpleNamespace(last_hidden_state=last_hidden_state, hidden_states=hidden_states) + def id_preprocess(x: str) -> str: return x @@ -74,8 +81,14 @@ def id_preprocess(x: str) -> str: def chat_list_preprocess(x: str): return [ - {"role": "system", "content": "Describe the video."}, - {"role": "user", "content": x if x else " "}, + { + "role": "system", + "content": "Describe the video." + }, + { + "role": "user", + "content": x if x else " " + }, ] @@ -83,6 +96,7 @@ def take_mean_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor: # [B, T, H] -> [B, H] return outputs.last_hidden_state.mean(dim=1) + def make_args(num_encoders=2, text_len=4, hidden_size=8): enc_cfgs = [] preprocess_fns = [] @@ -100,11 +114,13 @@ def make_args(num_encoders=2, text_len=4, hidden_size=8): ) return FastVideoArgs(model_path="", pipeline_config=pipe_cfg), hidden_size + def make_stage(num_encoders=2, hidden_size=8): tokenizers = [FakeTokenizer() for _ in range(num_encoders)] encoders = [FakeTextEncoder(hidden_size=hidden_size) for _ in range(num_encoders)] return TextEncodingStage(text_encoders=encoders, tokenizers=tokenizers) + def test_encode_text_selection_and_shapes(): fastvideo_args, hidden = make_args(num_encoders=2, text_len=4, hidden_size=8) stage = make_stage(num_encoders=2, hidden_size=hidden) @@ -122,19 +138,25 @@ def test_encode_text_selection_and_shapes(): assert masks2[0].shape == (1, 4) # dict return - d = stage.encode_text(["a","b"], fastvideo_args, encoder_index=[0,1], return_type="dict") + d = stage.encode_text(["a", "b"], fastvideo_args, encoder_index=[0, 1], return_type="dict") assert set(d.keys()) == {"0", "1"} assert d["0"].shape == (2, hidden) # stack return - s = stage.encode_text(["a","b"], fastvideo_args, encoder_index=[0,1], return_type="stack") + s = stage.encode_text(["a", "b"], fastvideo_args, encoder_index=[0, 1], return_type="stack") assert s.shape == (2, 2, hidden) # [encoders, batch, hidden] # overrides: dtype + max_length - e3, m3 = stage.encode_text(["a"], fastvideo_args, encoder_index=[0], dtype=torch.float16, return_attention_mask=True, max_length=3) + e3, m3 = stage.encode_text(["a"], + fastvideo_args, + encoder_index=[0], + dtype=torch.float16, + return_attention_mask=True, + max_length=3) assert e3[0].dtype == torch.float16 assert m3[0].shape[1] == 3 + def test_forward_integration_cfg_off_and_on(): fastvideo_args, hidden = make_args(num_encoders=2, text_len=4, hidden_size=8) stage = make_stage(num_encoders=2, hidden_size=hidden) @@ -222,8 +244,14 @@ def test_chat_list_preprocess_output_is_not_stripped(): assert embeds[0].shape == (1, hidden) assert masks[0].shape == (1, 5) assert tokenizer.last_messages == [[ - {"role": "system", "content": "Describe the video."}, - {"role": "user", "content": "a robotic arm welding a metal structure"}, + { + "role": "system", + "content": "Describe the video." + }, + { + "role": "user", + "content": "a robotic arm welding a metal structure" + }, ]] assert tokenizer.last_kwargs["return_tensors"] == "pt" @@ -337,18 +365,15 @@ class StaticBufferTextEncoder(torch.nn.Module): def __init__(self, text_len=4, hidden_size=8): super().__init__() - self.register_buffer("static_out", - torch.zeros(1, text_len, hidden_size)) - self.register_buffer("static_audio", - torch.zeros(1, text_len, hidden_size)) + self.register_buffer("static_out", torch.zeros(1, text_len, hidden_size)) + self.register_buffer("static_audio", torch.zeros(1, text_len, hidden_size)) def forward(self, input_ids, attention_mask, output_hidden_states=False): # A CUDAGraph replay rewrites the static output buffers in place. self.static_out += 1.0 self.static_audio += 2.0 hidden_states = (self.static_audio, ) if output_hidden_states else None - return types.SimpleNamespace(last_hidden_state=self.static_out, - hidden_states=hidden_states) + return types.SimpleNamespace(last_hidden_state=self.static_out, hidden_states=hidden_states) def identity_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor: @@ -371,11 +396,9 @@ def test_encode_text_output_does_not_alias_encoder_static_buffers(): cfg.output_hidden_states = True encoder = StaticBufferTextEncoder(text_len=4, hidden_size=hidden) - stage = TextEncodingStage(text_encoders=[encoder], - tokenizers=[FakeTokenizer()]) + stage = TextEncodingStage(text_encoders=[encoder], tokenizers=[FakeTokenizer()]) - embeds = stage.encode_text("a cat", fastvideo_args, encoder_index=[0], - device="cpu") + embeds = stage.encode_text("a cat", fastvideo_args, encoder_index=[0], device="cpu") prompt_embeds = embeds[0] audio_embeds = stage._last_audio_embeds[0] @@ -387,7 +410,6 @@ def test_encode_text_output_does_not_alias_encoder_static_buffers(): # that overwrites the encoder's static buffers in place. expected_prompt = prompt_embeds.clone() expected_audio = audio_embeds.clone() - stage.encode_text("bad quality", fastvideo_args, encoder_index=[0], - device="cpu") + stage.encode_text("bad quality", fastvideo_args, encoder_index=[0], device="cpu") torch.testing.assert_close(prompt_embeds, expected_prompt) torch.testing.assert_close(audio_embeds, expected_audio) diff --git a/fastvideo/tests/test_partial_hf_download.py b/fastvideo/tests/test_partial_hf_download.py index e7a6b6581..a2e393d90 100644 --- a/fastvideo/tests/test_partial_hf_download.py +++ b/fastvideo/tests/test_partial_hf_download.py @@ -147,11 +147,10 @@ def test_minimax_patterns_filter_unused_repo_partitions( "Ref2VA/model.safetensors", ] - selected = set( - filter_repo_objects( - repo_files, - allow_patterns=pipeline_cls.get_hf_download_allow_patterns(), - )) + selected = set(filter_repo_objects( + repo_files, + allow_patterns=pipeline_cls.get_hf_download_allow_patterns(), + )) assert "model_index.json" in selected assert "scheduler/scheduler_config.json" in selected @@ -234,10 +233,11 @@ def test_model_index_supports_umbrella_repo_paths(tmp_path: Path, monkeypatch) - captured.update(kwargs) manifest_path = tmp_path / kwargs["filename"] manifest_path.parent.mkdir(parents=True, exist_ok=True) - manifest_path.write_text(json.dumps({ - "_class_name": "MiniMaxH3ModularPipeline", - "_diffusers_version": "0.35.0", - })) + manifest_path.write_text( + json.dumps({ + "_class_name": "MiniMaxH3ModularPipeline", + "_diffusers_version": "0.35.0", + })) return str(manifest_path) monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_hf_hub_download) diff --git a/fastvideo/tests/train/callbacks/test_callback.py b/fastvideo/tests/train/callbacks/test_callback.py index 12c4b62c4..13a022a06 100644 --- a/fastvideo/tests/train/callbacks/test_callback.py +++ b/fastvideo/tests/train/callbacks/test_callback.py @@ -28,9 +28,7 @@ from fastvideo.train.callbacks.callback import ( # This helper attaches a temporary handler directly to the target # logger and yields the captured records. @contextmanager -def _capture_logger( - name: str, level: int = logging.WARNING -) -> Iterator[list[logging.LogRecord]]: +def _capture_logger(name: str, level: int = logging.WARNING) -> Iterator[list[logging.LogRecord]]: logger = logging.getLogger(name) records: list[logging.LogRecord] = [] @@ -65,9 +63,7 @@ class _RecordingCallback(Callback): def on_train_start(self, method, iteration: int = 0) -> None: self._sink.append(f"{self._tag}:on_train_start:{iteration}") - def on_training_step_end( - self, method, loss_dict, iteration: int = 0 - ) -> None: + def on_training_step_end(self, method, loss_dict, iteration: int = 0) -> None: self._sink.append(f"{self._tag}:on_training_step_end:{iteration}") def on_validation_begin(self, method, iteration: int = 0) -> None: @@ -97,9 +93,7 @@ class TestCallbackBase: def test_default_hooks_return_none(self) -> None: cb = Callback() assert cb.on_train_start(method=None) is None - assert ( - cb.on_training_step_end(method=None, loss_dict={}) is None - ) + assert (cb.on_training_step_end(method=None, loss_dict={}) is None) assert cb.on_before_optimizer_step(method=None) is None assert cb.on_validation_begin(method=None) is None assert cb.on_validation_end(method=None) is None @@ -131,8 +125,7 @@ class TestCallbackDictInit: assert "grad_clip" in cb_dict._callbacks from fastvideo.train.callbacks.grad_clip import ( - GradNormClipCallback, - ) + GradNormClipCallback, ) cb = cb_dict._callbacks["grad_clip"] assert isinstance(cb, GradNormClipCallback) # CallbackDict wires up training_config + back-pointer. @@ -142,10 +135,8 @@ class TestCallbackDictInit: def test_explicit_target_overrides_name_lookup(self) -> None: cfg = { "anything_goes": { - "_target_": ( - "fastvideo.train.callbacks.grad_clip." - "GradNormClipCallback" - ), + "_target_": ("fastvideo.train.callbacks.grad_clip." + "GradNormClipCallback"), "max_grad_norm": 1.0, } } @@ -154,23 +145,16 @@ class TestCallbackDictInit: def test_unknown_name_without_target_is_skipped(self) -> None: cfg = {"mystery": {"some_arg": 1}} - with _capture_logger( - "fastvideo.train.callbacks.callback" - ) as records: + with _capture_logger("fastvideo.train.callbacks.callback") as records: cb_dict = CallbackDict(cfg, training_config=object()) assert cb_dict._callbacks == {} - assert any( - "missing" in r.getMessage() and "mystery" in r.getMessage() - for r in records - ) + assert any("missing" in r.getMessage() and "mystery" in r.getMessage() for r in records) def test_non_callback_target_raises(self) -> None: cfg = { "bad": { - "_target_": ( - "fastvideo.tests.train.callbacks.test_callback." - "_NotACallback" - ) + "_target_": ("fastvideo.tests.train.callbacks.test_callback." + "_NotACallback") } } with pytest.raises(TypeError, match="expected a Callback"): @@ -192,16 +176,10 @@ class TestCallbackDictInit: class TestCallbackDictDispatch: - def _build( - self, sink: list[str] - ) -> CallbackDict: + def _build(self, sink: list[str]) -> CallbackDict: cb_dict = CallbackDict({}, training_config=object()) - cb_dict._callbacks["first"] = _RecordingCallback( - tag="first", sink=sink - ) - cb_dict._callbacks["second"] = _RecordingCallback( - tag="second", sink=sink - ) + cb_dict._callbacks["first"] = _RecordingCallback(tag="first", sink=sink) + cb_dict._callbacks["second"] = _RecordingCallback(tag="second", sink=sink) return cb_dict def test_dispatch_calls_all_in_insertion_order(self) -> None: @@ -217,14 +195,13 @@ class TestCallbackDictDispatch: def test_dispatch_to_hook_some_callbacks_skip(self) -> None: sink: list[str] = [] cb_dict = self._build(sink) + # The base Callback subclass below only implements one hook; # dispatch should still fan out without raising. class _OnlyValidation(Callback): - def on_validation_end( - self, method, iteration: int = 0 - ) -> None: + def on_validation_end(self, method, iteration: int = 0) -> None: sink.append(f"vend:{iteration}") cb_dict._callbacks["only_v"] = _OnlyValidation() @@ -252,12 +229,8 @@ class TestCallbackDictStateDict: def _build(self) -> tuple[CallbackDict, list[str]]: sink: list[str] = [] cb_dict = CallbackDict({}, training_config=object()) - cb_dict._callbacks["first"] = _RecordingCallback( - tag="first", sink=sink - ) - cb_dict._callbacks["second"] = _RecordingCallback( - tag="second", sink=sink - ) + cb_dict._callbacks["first"] = _RecordingCallback(tag="first", sink=sink) + cb_dict._callbacks["second"] = _RecordingCallback(tag="second", sink=sink) return cb_dict, sink def test_state_dict_returns_per_callback_dict(self) -> None: @@ -269,22 +242,19 @@ class TestCallbackDictStateDict: def test_load_state_dict_dispatches_to_each(self) -> None: cb_dict, sink = self._build() - cb_dict.load_state_dict( - { - "first": {"marker": 1}, - "second": {"marker": 2}, - } - ) + cb_dict.load_state_dict({ + "first": { + "marker": 1 + }, + "second": { + "marker": 2 + }, + }) assert sink == ["first:load:1", "second:load:2"] def test_load_state_dict_missing_key_warns_no_raise(self) -> None: cb_dict, sink = self._build() - with _capture_logger( - "fastvideo.train.callbacks.callback" - ) as records: + with _capture_logger("fastvideo.train.callbacks.callback") as records: cb_dict.load_state_dict({"first": {"marker": 99}}) assert sink == ["first:load:99"] - assert any( - "second" in r.getMessage() and "not found" in r.getMessage() - for r in records - ) + assert any("second" in r.getMessage() and "not found" in r.getMessage() for r in records) diff --git a/fastvideo/tests/train/callbacks/test_ema.py b/fastvideo/tests/train/callbacks/test_ema.py index bf10f1dd2..3e84079aa 100644 --- a/fastvideo/tests/train/callbacks/test_ema.py +++ b/fastvideo/tests/train/callbacks/test_ema.py @@ -16,7 +16,6 @@ import torch from fastvideo.train.callbacks.ema import EMACallback - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -89,9 +88,7 @@ class TestOnTrainingStepEnd: def test_no_op_before_train_start(self) -> None: cb = EMACallback() # student_ema is None until on_train_start. - cb.on_training_step_end( - _Method(transformer=None), loss_dict={}, iteration=0 - ) + cb.on_training_step_end(_Method(transformer=None), loss_dict={}, iteration=0) assert not cb._ema_started def test_skipped_until_start_iter(self) -> None: @@ -103,9 +100,7 @@ class TestOnTrainingStepEnd: with torch.no_grad(): transformer.weight.fill_(7.0) - cb.on_training_step_end( - _Method(transformer), loss_dict={}, iteration=5 - ) + cb.on_training_step_end(_Method(transformer), loss_dict={}, iteration=5) # Below start_iter: shadow is untouched, _ema_started False. assert not cb._ema_started assert torch.allclose( @@ -122,9 +117,7 @@ class TestOnTrainingStepEnd: with torch.no_grad(): transformer.weight.fill_(5.0) - cb.on_training_step_end( - _Method(transformer), loss_dict={}, iteration=10 - ) + cb.on_training_step_end(_Method(transformer), loss_dict={}, iteration=10) # First active step: shadow is re-initialized from the # current transformer (5.0) and *then* update() applies decay # against the same value, so shadow stays at 5.0. @@ -140,16 +133,12 @@ class TestOnTrainingStepEnd: cb.on_train_start(_Method(transformer), iteration=0) # Step 0: re-init at 2.0, then update against 2.0 → still 2.0. - cb.on_training_step_end( - _Method(transformer), loss_dict={}, iteration=0 - ) + cb.on_training_step_end(_Method(transformer), loss_dict={}, iteration=0) # Step 1: drift transformer to 12.0, expect # shadow = 0.9 * 2.0 + 0.1 * 12.0 = 3.0. with torch.no_grad(): transformer.weight.fill_(12.0) - cb.on_training_step_end( - _Method(transformer), loss_dict={}, iteration=1 - ) + cb.on_training_step_end(_Method(transformer), loss_dict={}, iteration=1) assert torch.allclose( cb.student_ema.shadow["weight"], torch.full((2, 4), 3.0), @@ -164,10 +153,7 @@ class TestOnTrainingStepEnd: cb.on_train_start(method, iteration=0) cb.on_training_step_end(method, loss_dict={}, iteration=0) - assert any( - payload.get("ema/decay") == 0.99 and step == 0 - for payload, step in tracker.entries - ) + assert any(payload.get("ema/decay") == 0.99 and step == 0 for payload, step in tracker.entries) # --------------------------------------------------------------------------- @@ -183,9 +169,7 @@ class TestEmaContext: # No on_train_start → student_ema is None. with cb.ema_context(transformer) as t: assert t is transformer - assert torch.allclose( - t.weight, torch.full((2, 4), 3.0) - ) + assert torch.allclose(t.weight, torch.full((2, 4), 3.0)) def test_swaps_weights_then_restores(self) -> None: transformer = _tiny_transformer(fill=1.0) @@ -203,9 +187,7 @@ class TestEmaContext: with cb.ema_context(transformer) as t: assert torch.allclose(t.weight, torch.full((2, 4), 1.0)) - assert torch.allclose( - transformer.weight, torch.full((2, 4), 9.0) - ) + assert torch.allclose(transformer.weight, torch.full((2, 4), 9.0)) # --------------------------------------------------------------------------- @@ -232,8 +214,7 @@ class TestStateDict: # Build a fresh callback and load. fresh = EMACallback(decay=0.5, start_iter=0) - fresh.on_train_start(_Method(_tiny_transformer(fill=0.0)), - iteration=0) + fresh.on_train_start(_Method(_tiny_transformer(fill=0.0)), iteration=0) # Sanity: fresh shadow != saved shadow before load. assert not torch.allclose( fresh.student_ema.shadow["weight"], diff --git a/fastvideo/tests/train/callbacks/test_grad_clip.py b/fastvideo/tests/train/callbacks/test_grad_clip.py index 1a68bf24c..3fee9f63e 100644 --- a/fastvideo/tests/train/callbacks/test_grad_clip.py +++ b/fastvideo/tests/train/callbacks/test_grad_clip.py @@ -12,7 +12,6 @@ import torch from fastvideo.train.callbacks.grad_clip import GradNormClipCallback - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -40,9 +39,7 @@ class _Method: self.tracker = tracker self.iter_seen: int | None = None - def get_grad_clip_targets( - self, iteration: int - ) -> dict[str, torch.nn.Module]: + def get_grad_clip_targets(self, iteration: int) -> dict[str, torch.nn.Module]: self.iter_seen = iteration return self._targets @@ -83,9 +80,7 @@ class TestGradNormClipCallback: assert _grad_norm(m) > 1.0 cb = GradNormClipCallback(max_grad_norm=1.0) - cb.on_before_optimizer_step( - method=_Method(targets={"m": m}), iteration=0 - ) + cb.on_before_optimizer_step(method=_Method(targets={"m": m}), iteration=0) # After clipping the L2 norm should not exceed max_grad_norm # (allow a tiny epsilon for the +1e-6 in the clip helper). @@ -97,9 +92,7 @@ class TestGradNormClipCallback: assert before < 1.0 cb = GradNormClipCallback(max_grad_norm=1.0) - cb.on_before_optimizer_step( - method=_Method(targets={"m": m}), iteration=0 - ) + cb.on_before_optimizer_step(method=_Method(targets={"m": m}), iteration=0) # Clip coef >1 is clamped to 1, so values are preserved # (modulo the *1.0 multiply, which is exact for floats). @@ -115,9 +108,7 @@ class TestGradNormClipCallback: def test_tracker_logged_when_enabled(self) -> None: m = _make_module(grad_value=5.0) tracker = _RecordingTracker() - cb = GradNormClipCallback( - max_grad_norm=1.0, log_grad_norms=True - ) + cb = GradNormClipCallback(max_grad_norm=1.0, log_grad_norms=True) cb.on_before_optimizer_step( method=_Method(targets={"layer": m}, tracker=tracker), iteration=7, @@ -131,9 +122,7 @@ class TestGradNormClipCallback: def test_tracker_not_logged_when_disabled(self) -> None: m = _make_module(grad_value=5.0) tracker = _RecordingTracker() - cb = GradNormClipCallback( - max_grad_norm=1.0, log_grad_norms=False - ) + cb = GradNormClipCallback(max_grad_norm=1.0, log_grad_norms=False) cb.on_before_optimizer_step( method=_Method(targets={"m": m}, tracker=tracker), iteration=0, @@ -143,13 +132,12 @@ class TestGradNormClipCallback: def test_no_tracker_does_not_raise(self) -> None: m = _make_module(grad_value=5.0) cb = GradNormClipCallback(max_grad_norm=1.0, log_grad_norms=True) + # Method without a tracker attribute at all. class _BareMethod: - def get_grad_clip_targets( - self, iteration: int - ) -> dict[str, torch.nn.Module]: + def get_grad_clip_targets(self, iteration: int) -> dict[str, torch.nn.Module]: return {"m": m} cb.on_before_optimizer_step(method=_BareMethod(), iteration=0) diff --git a/fastvideo/tests/train/callbacks/test_validation.py b/fastvideo/tests/train/callbacks/test_validation.py index 015625938..e8baf8d1d 100644 --- a/fastvideo/tests/train/callbacks/test_validation.py +++ b/fastvideo/tests/train/callbacks/test_validation.py @@ -31,12 +31,10 @@ from fastvideo.train.callbacks.validation import ( _ValidationMetricStats, ) - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- - _PIPE_TARGET = "fastvideo.pipelines.basic.wan.wan_pipeline.WanPipeline" @@ -276,11 +274,8 @@ class TestH3ValidationContract: num_width=96, num_latent_t=2, ), - pipeline_config=SimpleNamespace( - vae_config=SimpleNamespace( - arch_config=SimpleNamespace(temporal_compression_ratio=4), - ), - ), + pipeline_config=SimpleNamespace(vae_config=SimpleNamespace( + arch_config=SimpleNamespace(temporal_compression_ratio=4), ), ), model_path="unused", vsa_sparsity=0.0, ) @@ -312,11 +307,8 @@ class TestH3ValidationContract: num_width=96, num_latent_t=2, ), - pipeline_config=SimpleNamespace( - vae_config=SimpleNamespace( - arch_config=SimpleNamespace(temporal_compression_ratio=4), - ), - ), + pipeline_config=SimpleNamespace(vae_config=SimpleNamespace( + arch_config=SimpleNamespace(temporal_compression_ratio=4), ), ), model_path="unused", vsa_sparsity=0.0, ) @@ -382,9 +374,8 @@ class TestH3ValidationContract: monkeypatch.setattr( cb, "_prepare_validation_batch", - lambda sampling_param, validation_batch, num_inference_steps: SimpleNamespace( - prompt=validation_batch["caption"], - ), + lambda sampling_param, validation_batch, num_inference_steps: SimpleNamespace(prompt=validation_batch[ + "caption"], ), ) result = cb._run_validation_for_steps(50, transformer=torch.nn.Identity()) @@ -399,6 +390,7 @@ class TestH3ValidationContract: def test_log_validation_video_artifacts_emits_eight_wandb_videos(self) -> None: """Verify a MiniMax H3 event contains eight media objects and scalar evidence.""" + class FakeWandbTracker: """Record tracker calls without contacting W&B.""" @@ -547,9 +539,7 @@ class TestStateDict: # Receiver: fresh generator with a different seed. fresh = _make_callback() - fresh.validation_random_generator = ( - torch.Generator(device="cpu").manual_seed(999) - ) + fresh.validation_random_generator = (torch.Generator(device="cpu").manual_seed(999)) fresh.load_state_dict(state) # After load, both generators draw the same next sample. @@ -561,9 +551,7 @@ class TestStateDict: cb = _make_callback() # Generator is None: load must not raise even when state has # an rng entry. - cb.load_state_dict( - {"validation_rng": torch.tensor([1, 2, 3], dtype=torch.uint8)} - ) + cb.load_state_dict({"validation_rng": torch.tensor([1, 2, 3], dtype=torch.uint8)}) assert cb.validation_random_generator is None @@ -581,16 +569,21 @@ class TestMetricAggregation: stats, row, { - "vbench.aesthetic_quality": SimpleNamespace( + "vbench.aesthetic_quality": + SimpleNamespace( name="vbench.aesthetic_quality", score=0.5, details={"ignored": [1, 2, 3]}, ), - SYNTHETIC_OPTICAL_FLOW_METRIC: SimpleNamespace( + SYNTHETIC_OPTICAL_FLOW_METRIC: + SimpleNamespace( name=SYNTHETIC_OPTICAL_FLOW_METRIC, score=1.5, details={ - **{key: float(i) for i, key in enumerate(SYNTHETIC_OPTICAL_FLOW_LOG_KEYS)}, + **{ + key: float(i) + for i, key in enumerate(SYNTHETIC_OPTICAL_FLOW_LOG_KEYS) + }, "pixel_epe_mean_std": 99.0, "pixel_epe_mean_max": 100.0, "pixel_epe_mean_auc": 101.0, @@ -616,13 +609,23 @@ class TestMetricAggregation: dst = _ValidationMetricStats( sums={"a": 1.0}, counts={"a": 1.0}, - per_video=[{"path": "a.mp4"}], + per_video=[{ + "path": "a.mp4" + }], errors=["first"], ) src = _ValidationMetricStats( - sums={"a": 2.0, "b": 3.0}, - counts={"a": 2.0, "b": 1.0}, - per_video=[{"path": "b.mp4"}], + sums={ + "a": 2.0, + "b": 3.0 + }, + counts={ + "a": 2.0, + "b": 1.0 + }, + per_video=[{ + "path": "b.mp4" + }], errors=["second"], ) @@ -640,10 +643,7 @@ class TestMetricAggregation: metrics={"names": ["vbench.aesthetic_quality"]}, ) - assert ( - cb._metric_log_name("vbench.aesthetic_quality") - == "metrics/validation/vbench/aesthetic_quality" - ) + assert (cb._metric_log_name("vbench.aesthetic_quality") == "metrics/validation/vbench/aesthetic_quality") def test_metric_device_uses_local_rank( self, @@ -694,12 +694,10 @@ class TestMetricAggregation: keyboard = np.zeros((4, 6)) mouse = np.zeros((4, 2)) - actions = ValidationCallback._validation_actions( - { - "keyboard_cond": keyboard, - "mouse_cond": mouse, - } - ) + actions = ValidationCallback._validation_actions({ + "keyboard_cond": keyboard, + "mouse_cond": mouse, + }) assert actions is not None assert np.array_equal(actions["keyboard"], keyboard) @@ -773,9 +771,7 @@ class TestActionOverlay: suffix="_overlay", ) - assert saved.filenames == [ - str(tmp_path / "validation_step_7_inference_steps_4_rank_3_video_0_overlay.mp4") - ] + assert saved.filenames == [str(tmp_path / "validation_step_7_inference_steps_4_rank_3_video_0_overlay.mp4")] assert saved.indices == [0] assert calls == [(saved.filenames[0], 25)] @@ -820,9 +816,7 @@ class TestActionOverlay: ) assert saved.indices == [1] - assert saved.filenames == [ - str(tmp_path / "validation_step_8_inference_steps_5_rank_0_video_1.mp4") - ] + assert saved.filenames == [str(tmp_path / "validation_step_8_inference_steps_5_rank_0_video_1.mp4")] assert len(calls) == 2 def test_save_validation_videos_counts_preserved_audio( diff --git a/fastvideo/tests/train/methods/grad_norm_regression.py b/fastvideo/tests/train/methods/grad_norm_regression.py index fda6a3bb5..f579108a5 100644 --- a/fastvideo/tests/train/methods/grad_norm_regression.py +++ b/fastvideo/tests/train/methods/grad_norm_regression.py @@ -111,13 +111,9 @@ def layer0_grad_norm(transformer) -> float: (``.item()``) at the end, rather than one per parameter. """ blocks = resolve_blocks(transformer) - assert blocks is not None and len(blocks) > 0, ( - "transformer is expected to expose a non-empty block list " - f"(one of {_BLOCK_LIST_ATTRS})") - grads = [ - p.grad for p in blocks[0].parameters() - if p.requires_grad and p.grad is not None - ] + assert blocks is not None and len(blocks) > 0, ("transformer is expected to expose a non-empty block list " + f"(one of {_BLOCK_LIST_ATTRS})") + grads = [p.grad for p in blocks[0].parameters() if p.requires_grad and p.grad is not None] if not grads: return 0.0 # Some loaders (e.g. Cosmos via fsdp_load) keep parameters/grads as @@ -137,9 +133,7 @@ def _load_refs() -> dict[str, dict[str, float]]: def _save_refs(refs: dict[str, dict[str, float]]) -> None: - _REFS_PATH.write_text( - json.dumps(refs, indent=2, sort_keys=True) + "\n", - encoding="utf-8") + _REFS_PATH.write_text(json.dumps(refs, indent=2, sort_keys=True) + "\n", encoding="utf-8") def check_grad_norm_regression( @@ -160,26 +154,22 @@ def check_grad_norm_regression( if os.environ.get(_UPDATE_ENV) == "1": if device_key is None: - pytest.skip( - f"{_UPDATE_ENV}=1 but GPU '{_device_name()}' has no reference " - "key; add it to _DEVICE_MAPPINGS first") + pytest.skip(f"{_UPDATE_ENV}=1 but GPU '{_device_name()}' has no reference " + "key; add it to _DEVICE_MAPPINGS first") refs = _load_refs() refs.setdefault(test_name, {})[device_key] = round(norm, 4) _save_refs(refs) - pytest.skip( - f"recorded grad-norm reference {test_name}[{device_key}] = " - f"{norm:.4f} (assertion skipped under {_UPDATE_ENV}=1)") + pytest.skip(f"recorded grad-norm reference {test_name}[{device_key}] = " + f"{norm:.4f} (assertion skipped under {_UPDATE_ENV}=1)") ref = _load_refs().get(test_name, {}).get(device_key) \ if device_key is not None else None if ref is None: - pytest.skip( - f"no grad-norm reference for {test_name} on '{_device_name()}' " - f"(device_key={device_key}); run with {_UPDATE_ENV}=1 to seed it") + pytest.skip(f"no grad-norm reference for {test_name} on '{_device_name()}' " + f"(device_key={device_key}); run with {_UPDATE_ENV}=1 to seed it") rel = abs(norm - ref) / (abs(ref) + 1e-12) - assert rel <= rtol, ( - f"{test_name}[{device_key}] grad-norm regression: got {norm:.4f}, " - f"reference {ref:.4f}, relative error {rel:.3%} exceeds rtol " - f"{rtol:.0%}. If this is an intentional change, refresh the reference " - f"with {_UPDATE_ENV}=1 and explain why in the PR.") + assert rel <= rtol, (f"{test_name}[{device_key}] grad-norm regression: got {norm:.4f}, " + f"reference {ref:.4f}, relative error {rel:.3%} exceeds rtol " + f"{rtol:.0%}. If this is an intentional change, refresh the reference " + f"with {_UPDATE_ENV}=1 and explain why in the PR.") diff --git a/fastvideo/tests/train/methods/test_cosmos_finetune.py b/fastvideo/tests/train/methods/test_cosmos_finetune.py index edfa83d3b..3029a4567 100644 --- a/fastvideo/tests/train/methods/test_cosmos_finetune.py +++ b/fastvideo/tests/train/methods/test_cosmos_finetune.py @@ -32,9 +32,7 @@ from .grad_norm_regression import ( resolve_blocks, ) -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" - / "cosmos_t2w_finetune_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "cosmos_t2w_finetune_min.yaml") # Cosmos 2.5 Reason1 (Qwen2.5-VL) text embedding width. _COSMOS_TEXT_DIM = 100352 @@ -51,19 +49,14 @@ def _build_synthetic_batch( """ batch_size = 1 return { - "text_embedding": - torch.randn(batch_size, 16, _COSMOS_TEXT_DIM, device=device, - dtype=dtype), - "text_attention_mask": - torch.ones(batch_size, 16, device=device), - "vae_latent": - torch.randn(batch_size, 16, 4, 8, 8, device=device, dtype=dtype), + "text_embedding": torch.randn(batch_size, 16, _COSMOS_TEXT_DIM, device=device, dtype=dtype), + "text_attention_mask": torch.ones(batch_size, 16, device=device), + "vae_latent": torch.randn(batch_size, 16, 4, 8, 8, device=device, dtype=dtype), } @pytest.mark.usefixtures("distributed_setup") -def test_cosmos_finetune_single_train_step( - monkeypatch: pytest.MonkeyPatch) -> None: +def test_cosmos_finetune_single_train_step(monkeypatch: pytest.MonkeyPatch) -> None: if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -100,14 +93,12 @@ def test_cosmos_finetune_single_train_step( loss = loss_map["total_loss"] assert torch.is_tensor(loss), "total_loss must be a torch.Tensor" - assert torch.isfinite(loss).item(), ( - f"total_loss is not finite: {loss.item()}") + assert torch.isfinite(loss).item(), (f"total_loss is not finite: {loss.item()}") method.backward(loss_map, outputs, grad_accum_rounds=1) blocks = resolve_blocks(model.transformer) - assert blocks is not None and len(blocks) > 0, ( - "transformer is expected to expose a non-empty block list") + assert blocks is not None and len(blocks) > 0, ("transformer is expected to expose a non-empty block list") layer0 = blocks[0] trainable = [p for p in layer0.parameters() if p.requires_grad] @@ -115,14 +106,11 @@ def test_cosmos_finetune_single_train_step( for i, p in enumerate(trainable): assert p.grad is not None, f"layer 0 param[{i}] has None grad" - assert torch.isfinite(p.grad).all().item(), ( - f"layer 0 param[{i}] grad contains NaN/Inf") + assert torch.isfinite(p.grad).all().item(), (f"layer 0 param[{i}] grad contains NaN/Inf") - any_nonzero = any( - p.grad.detach().float().norm().item() > 0.0 for p in trainable) - assert any_nonzero, ( - "all layer-0 grads are exactly zero; backward did not " - "reach the first transformer block") + any_nonzero = any(p.grad.detach().float().norm().item() > 0.0 for p in trainable) + assert any_nonzero, ("all layer-0 grads are exactly zero; backward did not " + "reach the first transformer block") # 5a-ii: device-keyed grad-norm regression on top of the same harness. # Skips when the current GPU has no seeded reference. diff --git a/fastvideo/tests/train/methods/test_ltx2_finetune.py b/fastvideo/tests/train/methods/test_ltx2_finetune.py index 31a1bad39..5a9bc3467 100644 --- a/fastvideo/tests/train/methods/test_ltx2_finetune.py +++ b/fastvideo/tests/train/methods/test_ltx2_finetune.py @@ -62,9 +62,7 @@ def _gpu_too_small() -> bool: def _flashinfer_fp4_available() -> bool: try: from flashinfer import ( # noqa: F401 - SfLayout, - mm_fp4, - nvfp4_quantize, + SfLayout, mm_fp4, nvfp4_quantize, ) except ImportError: return False @@ -83,23 +81,15 @@ def _build_synthetic_batch( """ batch_size = 1 return { - "text_embedding": - torch.randn(batch_size, - _LTX2_TEXT_LEN, - text_dim, - device=device, - dtype=dtype), - "text_attention_mask": - torch.ones(batch_size, _LTX2_TEXT_LEN, device=device), - "vae_latent": - torch.randn(batch_size, 128, 3, 8, 8, device=device, dtype=dtype), + "text_embedding": torch.randn(batch_size, _LTX2_TEXT_LEN, text_dim, device=device, dtype=dtype), + "text_attention_mask": torch.ones(batch_size, _LTX2_TEXT_LEN, device=device), + "vae_latent": torch.randn(batch_size, 128, 3, 8, 8, device=device, dtype=dtype), } @pytest.mark.usefixtures("distributed_setup") @pytest.mark.parametrize("case", _CASES.keys()) -def test_ltx2_finetune_single_train_step( - case: str, monkeypatch: pytest.MonkeyPatch) -> None: +def test_ltx2_finetune_single_train_step(case: str, monkeypatch: pytest.MonkeyPatch) -> None: if _gpu_too_small(): pytest.skip(f"requires a CUDA GPU with >= {_MIN_GPU_MEMORY_GB}GB " "memory (LTX-2 DiT is 18.9B params)") @@ -145,21 +135,17 @@ def test_ltx2_finetune_single_train_step( from fastvideo.layers.quantization.nvfp4_qat_train_config import ( NVFP4QATTrainQuantizeMethod, ) quantized = sum( - isinstance(getattr(module, "quant_method", None), - NVFP4QATTrainQuantizeMethod) + isinstance(getattr(module, "quant_method", None), NVFP4QATTrainQuantizeMethod) for module in model.transformer.modules()) - assert quantized == expected_qat_linears, ( - f"expected {expected_qat_linears} deployment-matched NVFP4-QAT " - f"linears on the LTX-2 DiT, found {quantized}") + assert quantized == expected_qat_linears, (f"expected {expected_qat_linears} deployment-matched NVFP4-QAT " + f"linears on the LTX-2 DiT, found {quantized}") from fastvideo.attention.backends.attn_qat_train import ( AttnQatTrainImpl, ) qat_attention = sum( - isinstance(getattr(module, "attn_impl", None), AttnQatTrainImpl) - for module in model.transformer.modules()) - assert qat_attention == 96, ( - "expected ATTN_QAT_TRAIN on attn1 and attn2 in all 48 LTX-2 " - f"blocks, found {qat_attention} implementations") + isinstance(getattr(module, "attn_impl", None), AttnQatTrainImpl) for module in model.transformer.modules()) + assert qat_attention == 96, ("expected ATTN_QAT_TRAIN on attn1 and attn2 in all 48 LTX-2 " + f"blocks, found {qat_attention} implementations") method = FineTuneMethod( cfg=cfg, @@ -172,15 +158,13 @@ def test_ltx2_finetune_single_train_step( loss = loss_map["total_loss"] assert torch.is_tensor(loss), "total_loss must be a torch.Tensor" - assert torch.isfinite(loss).item(), ( - f"total_loss is not finite: {loss.item()}") + assert torch.isfinite(loss).item(), (f"total_loss is not finite: {loss.item()}") method.backward(loss_map, outputs, grad_accum_rounds=1) # LTX-2 nests its transformer_blocks under the ``model`` submodule. blocks = resolve_blocks(model.transformer.model) - assert blocks is not None and len(blocks) > 0, ( - "transformer is expected to expose a non-empty block list") + assert blocks is not None and len(blocks) > 0, ("transformer is expected to expose a non-empty block list") layer0 = blocks[0] trainable = [p for p in layer0.parameters() if p.requires_grad] @@ -188,24 +172,19 @@ def test_ltx2_finetune_single_train_step( for i, p in enumerate(trainable): assert p.grad is not None, f"layer 0 param[{i}] has None grad" - assert torch.isfinite(p.grad).all().item(), ( - f"layer 0 param[{i}] grad contains NaN/Inf") + assert torch.isfinite(p.grad).all().item(), (f"layer 0 param[{i}] grad contains NaN/Inf") - any_nonzero = any( - p.grad.detach().float().norm().item() > 0.0 for p in trainable) - assert any_nonzero, ( - "all layer-0 grads are exactly zero; backward did not " - "reach the first transformer block") + any_nonzero = any(p.grad.detach().float().norm().item() > 0.0 for p in trainable) + assert any_nonzero, ("all layer-0 grads are exactly zero; backward did not " + "reach the first transformer block") # Audio / cross-modal parameters must be frozen by default. audio_trainable = [ name for name, param in model.transformer.named_parameters() - if param.requires_grad and any( - pattern in name for pattern in ("audio", "a2v", "v2a", "av_ca")) + if param.requires_grad and any(pattern in name for pattern in ("audio", "a2v", "v2a", "av_ca")) ] - assert not audio_trainable, ( - f"audio/cross-modal params unexpectedly trainable: " - f"{audio_trainable[:5]}") + assert not audio_trainable, (f"audio/cross-modal params unexpectedly trainable: " + f"{audio_trainable[:5]}") # Device-keyed grad-norm regression on top of the same harness. # Skips when the current GPU has no seeded reference. diff --git a/fastvideo/tests/train/methods/test_matrixgame2_finetune.py b/fastvideo/tests/train/methods/test_matrixgame2_finetune.py index 6290673ed..dc32afa80 100644 --- a/fastvideo/tests/train/methods/test_matrixgame2_finetune.py +++ b/fastvideo/tests/train/methods/test_matrixgame2_finetune.py @@ -39,9 +39,7 @@ from .grad_norm_regression import ( resolve_blocks, ) -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" - / "matrixgame2_finetune_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "matrixgame2_finetune_min.yaml") # Matrix-Game 2.0 CLIP image-embedding width. _MG2_IMAGE_DIM = 1280 @@ -69,32 +67,24 @@ def _build_synthetic_batch( # action-module params receive gradients. action_frames = 13 return { - "vae_latent": - torch.randn(batch_size, 16, 4, 8, 8, device=device, dtype=dtype), - "clip_feature": - torch.randn(batch_size, 257, _MG2_IMAGE_DIM, device=device, - dtype=dtype), - "first_frame_latent": - torch.randn(batch_size, 16, 4, 8, 8, device=device, dtype=dtype), - "keyboard_cond": - torch.ones(batch_size, action_frames, 6, device=device, dtype=dtype), - "mouse_cond": - torch.ones(batch_size, action_frames, 2, device=device, dtype=dtype), + "vae_latent": torch.randn(batch_size, 16, 4, 8, 8, device=device, dtype=dtype), + "clip_feature": torch.randn(batch_size, 257, _MG2_IMAGE_DIM, device=device, dtype=dtype), + "first_frame_latent": torch.randn(batch_size, 16, 4, 8, 8, device=device, dtype=dtype), + "keyboard_cond": torch.ones(batch_size, action_frames, 6, device=device, dtype=dtype), + "mouse_cond": torch.ones(batch_size, action_frames, 2, device=device, dtype=dtype), } @pytest.mark.usefixtures("distributed_setup") -def test_matrixgame2_finetune_single_train_step( - monkeypatch: pytest.MonkeyPatch) -> None: +def test_matrixgame2_finetune_single_train_step(monkeypatch: pytest.MonkeyPatch) -> None: if not torch.cuda.is_available(): pytest.skip("requires CUDA") total_mem = torch.cuda.get_device_properties(0).total_memory if total_mem < _MIN_GPU_MEM_BYTES: - pytest.skip( - f"Matrix-Game 2.0 (~14B) backward needs ~56 GB; this GPU has " - f"{total_mem / 1024**3:.0f} GB. Loading is covered by " - "test_load_matrixgame2.py; the grad path runs on larger dev GPUs.") + pytest.skip(f"Matrix-Game 2.0 (~14B) backward needs ~56 GB; this GPU has " + f"{total_mem / 1024**3:.0f} GB. Loading is covered by " + "test_load_matrixgame2.py; the grad path runs on larger dev GPUs.") cfg = load_run_config(_FIXTURE) @@ -127,14 +117,12 @@ def test_matrixgame2_finetune_single_train_step( loss = loss_map["total_loss"] assert torch.is_tensor(loss), "total_loss must be a torch.Tensor" - assert torch.isfinite(loss).item(), ( - f"total_loss is not finite: {loss.item()}") + assert torch.isfinite(loss).item(), (f"total_loss is not finite: {loss.item()}") method.backward(loss_map, outputs, grad_accum_rounds=1) blocks = resolve_blocks(model.transformer) - assert blocks is not None and len(blocks) > 0, ( - "transformer is expected to expose a non-empty block list") + assert blocks is not None and len(blocks) > 0, ("transformer is expected to expose a non-empty block list") layer0 = blocks[0] named = [(n, p) for n, p in layer0.named_parameters() if p.requires_grad] @@ -146,19 +134,15 @@ def test_matrixgame2_finetune_single_train_step( # *active* path (self-attn, ffn, image cross-attn, action modules) is # intact rather than requiring every block-0 param to receive a grad. with_grad = [(n, p) for n, p in named if p.grad is not None] - assert len(with_grad) > 0, ( - "no layer-0 params received a gradient; backward did not reach " - "the first transformer block") + assert len(with_grad) > 0, ("no layer-0 params received a gradient; backward did not reach " + "the first transformer block") for n, p in with_grad: - assert torch.isfinite(p.grad).all().item(), ( - f"layer 0 param '{n}' grad contains NaN/Inf") + assert torch.isfinite(p.grad).all().item(), (f"layer 0 param '{n}' grad contains NaN/Inf") - any_nonzero = any( - p.grad.detach().float().norm().item() > 0.0 for _, p in with_grad) - assert any_nonzero, ( - "all layer-0 grads are exactly zero; backward did not " - "reach the first transformer block") + any_nonzero = any(p.grad.detach().float().norm().item() > 0.0 for _, p in with_grad) + assert any_nonzero, ("all layer-0 grads are exactly zero; backward did not " + "reach the first transformer block") # 5a-ii: device-keyed grad-norm regression on top of the same harness. # Skips when the current GPU has no seeded reference. diff --git a/fastvideo/tests/train/methods/test_minimax_h3_finetune.py b/fastvideo/tests/train/methods/test_minimax_h3_finetune.py index 7f4b435cf..0d1663e1a 100644 --- a/fastvideo/tests/train/methods/test_minimax_h3_finetune.py +++ b/fastvideo/tests/train/methods/test_minimax_h3_finetune.py @@ -29,23 +29,15 @@ from fastvideo.train.utils.config import load_run_config _FIXTURE = Path(__file__).resolve().parent.parent / "fixtures" / "minimax_h3_t2va_min.yaml" _REPO_ROOT = Path(__file__).resolve().parents[4] _EXPERIMENT_CONFIG = _REPO_ROOT / "examples/train/configs/overfit_minimax_h3_t2va.yaml" -_VALIDATION_PROMPTS = ( - _REPO_ROOT / "examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" -) +_VALIDATION_PROMPTS = (_REPO_ROOT / "examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json") _EXPECTED_VALIDATION_CAPTIONS = [ - ( - "A large metal cylinder is seen pressing down on a pile of Oreo cookies, " - "flattening them as if they were under a hydraulic press." - ), - ( - "A large metal cylinder is seen compressing colorful clay into a compact shape, " - "demonstrating the power of a hydraulic press." - ), - ( - "A large metal cylinder is seen pressing down on a pile of colorful candies, " - "flattening them as if they were under a hydraulic press. The candies are crushed " - "and broken into small pieces, creating a mess on the table." - ), + ("A large metal cylinder is seen pressing down on a pile of Oreo cookies, " + "flattening them as if they were under a hydraulic press."), + ("A large metal cylinder is seen compressing colorful clay into a compact shape, " + "demonstrating the power of a hydraulic press."), + ("A large metal cylinder is seen pressing down on a pile of colorful candies, " + "flattening them as if they were under a hydraulic press. The candies are crushed " + "and broken into small pieces, creating a mess on the table."), ] @@ -246,11 +238,8 @@ def test_h3_experiment_config_uses_modular_validation_callback() -> None: assert config["method"]["_target_"] == "fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod" assert validation["_target_"] == "fastvideo.train.callbacks.validation.ValidationCallback" assert validation["pipeline_target"] == ( - "fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline" - ) - assert validation["dataset_file"] == ( - "examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" - ) + "fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline") + assert validation["dataset_file"] == ("examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json") assert validation["every_steps"] == 20 assert validation["run_at_start"] is True assert validation["sampling_steps"] == [50] @@ -279,17 +268,12 @@ def test_h3_experiment_config_resolves_64_gpu_mesh_and_global_batch() -> None: "pin_cpu_memory": True, } data_parallel_degree = distributed["num_gpus"] // distributed["sp_size"] - effective_global_batch = ( - training["data"]["train_batch_size"] - * data_parallel_degree - * training["loop"]["gradient_accumulation_steps"] - ) + effective_global_batch = (training["data"]["train_batch_size"] * data_parallel_degree * + training["loop"]["gradient_accumulation_steps"]) assert distributed["hsdp_replicate_dim"] * distributed["hsdp_shard_dim"] == 64 assert data_parallel_degree == 8 assert effective_global_batch == 8 - assert training["data"]["data_path"] == { - "data/crush-smol_h3_t2va_single_sample_preprocessed": 8 - } + assert training["data"]["data_path"] == {"data/crush-smol_h3_t2va_single_sample_preprocessed": 8} assert training["loop"]["max_train_steps"] == 400 assert training["checkpoint"]["training_state_checkpointing_steps"] == 20 assert training["checkpoint"]["checkpoints_total_limit"] == 2 @@ -335,9 +319,7 @@ def test_h3_experiment_data_repeat_supplies_every_sp_group( assert sorted(assigned_indices) == list(range(8)) -def test_h3_slurm_scripts_preserve_inherited_wandb_key( - tmp_path: Path, -) -> None: +def test_h3_slurm_scripts_preserve_inherited_wandb_key(tmp_path: Path, ) -> None: """Verify shell syntax, H200 allocation, and secret-free job-script text.""" run_slurm = _REPO_ROOT / "examples/train/run_slurm.sh" launcher = _REPO_ROOT / "examples/train/launch_minimax_h3_t2va_crush_smol_validation.sh" @@ -350,12 +332,10 @@ def test_h3_slurm_scripts_preserve_inherited_wandb_key( generated_job_script = tmp_path / "generated-job.sh" sbatch_arguments = tmp_path / "sbatch-arguments.txt" fake_sbatch = fake_bin / "sbatch" - fake_sbatch.write_text( - "#!/bin/bash\n" - "printf '%s\\n' \"$@\" > \"$SBATCH_ARGUMENTS\"\n" - "dd status=none of=\"$SBATCH_JOB_SCRIPT\"\n" - "printf 'Submitted batch job 4242\\n'\n" - ) + fake_sbatch.write_text("#!/bin/bash\n" + "printf '%s\\n' \"$@\" > \"$SBATCH_ARGUMENTS\"\n" + "dd status=none of=\"$SBATCH_JOB_SCRIPT\"\n" + "printf 'Submitted batch job 4242\\n'\n") fake_sbatch.chmod(0o755) secret_value = "wandb-test-secret-must-not-appear" environment = os.environ.copy() @@ -403,30 +383,24 @@ def test_h3_slurm_job_resolves_one_node_rank_per_srun_task(tmp_path: Path) -> No torchrun_arguments_dir.mkdir() fake_sbatch = fake_bin / "sbatch" - fake_sbatch.write_text( - "#!/bin/bash\n" - "dd status=none of=\"$SBATCH_JOB_SCRIPT\"\n" - "printf 'Submitted batch job 4242\\n'\n" - ) + fake_sbatch.write_text("#!/bin/bash\n" + "dd status=none of=\"$SBATCH_JOB_SCRIPT\"\n" + "printf 'Submitted batch job 4242\\n'\n") fake_sbatch.chmod(0o755) fake_scontrol = fake_bin / "scontrol" fake_scontrol.write_text("#!/bin/bash\nprintf 'node%s\\n' {0..7}\n") fake_scontrol.chmod(0o755) fake_srun = fake_bin / "srun" - fake_srun.write_text( - "#!/bin/bash\n" - "set -euo pipefail\n" - "for rank in {0..7}; do\n" - " SLURM_PROCID=\"$rank\" \"$@\"\n" - "done\n" - ) + fake_srun.write_text("#!/bin/bash\n" + "set -euo pipefail\n" + "for rank in {0..7}; do\n" + " SLURM_PROCID=\"$rank\" \"$@\"\n" + "done\n") fake_srun.chmod(0o755) fake_torchrun = fake_bin / "torchrun" - fake_torchrun.write_text( - "#!/bin/bash\n" - "printf '%s\\n' \"$@\" > " - "\"$TORCHRUN_ARGUMENTS_DIR/rank_${SLURM_PROCID}.txt\"\n" - ) + fake_torchrun.write_text("#!/bin/bash\n" + "printf '%s\\n' \"$@\" > " + "\"$TORCHRUN_ARGUMENTS_DIR/rank_${SLURM_PROCID}.txt\"\n") fake_torchrun.chmod(0o755) environment = os.environ.copy() diff --git a/fastvideo/tests/train/methods/test_streaming_long_tuning.py b/fastvideo/tests/train/methods/test_streaming_long_tuning.py index 8a9c04190..d58a36bbc 100644 --- a/fastvideo/tests/train/methods/test_streaming_long_tuning.py +++ b/fastvideo/tests/train/methods/test_streaming_long_tuning.py @@ -38,9 +38,7 @@ from fastvideo.train.utils.config import load_run_config from .grad_norm_regression import check_grad_norm_regression -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" - / "wan_causal_t2v_streaming_long_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "wan_causal_t2v_streaming_long_min.yaml") def _build_synthetic_batch( @@ -49,12 +47,9 @@ def _build_synthetic_batch( ) -> dict[str, torch.Tensor]: batch_size = 1 return { - "text_embedding": - torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), - "text_attention_mask": - torch.ones(batch_size, 16, device=device), - "vae_latent": - torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype), + "text_embedding": torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), + "text_attention_mask": torch.ones(batch_size, 16, device=device), + "vae_latent": torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype), } @@ -65,25 +60,21 @@ def _assert_layer0_grads(model, *, role: str) -> None: assert len(trainable) > 0, f"{role} layer 0 has no trainable parameters" for i, p in enumerate(trainable): assert p.grad is not None, f"{role} layer 0 param[{i}] has None grad" - assert torch.isfinite(p.grad).all().item(), ( - f"{role} layer 0 param[{i}] grad contains NaN/Inf") - assert any( - p.grad.detach().float().norm().item() > 0.0 for p in trainable), ( - f"all {role} layer-0 grads are exactly zero; the loss did not " - "reach the first transformer block") + assert torch.isfinite(p.grad).all().item(), (f"{role} layer 0 param[{i}] grad contains NaN/Inf") + assert any(p.grad.detach().float().norm().item() > 0.0 + for p in trainable), (f"all {role} layer-0 grads are exactly zero; the loss did not " + "reach the first transformer block") def _assert_finite_losses(loss_map: dict[str, torch.Tensor]) -> None: for name in ("total_loss", "generator_loss", "fake_score_loss"): loss = loss_map[name] assert torch.is_tensor(loss), f"{name} must be a torch.Tensor" - assert torch.isfinite(loss).item(), ( - f"{name} is not finite: {loss.item()}") + assert torch.isfinite(loss).item(), (f"{name} is not finite: {loss.item()}") @pytest.mark.usefixtures("distributed_setup") -def test_streaming_long_tuning_multi_stage_train_steps( - monkeypatch: pytest.MonkeyPatch) -> None: +def test_streaming_long_tuning_multi_stage_train_steps(monkeypatch: pytest.MonkeyPatch) -> None: if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -140,9 +131,7 @@ def test_streaming_long_tuning_multi_stage_train_steps( method.backward(loss_map, outputs, grad_accum_rounds=1) _assert_layer0_grads(student, role="student") _assert_layer0_grads(critic, role="critic") - assert all(not p.requires_grad - for p in teacher.transformer.parameters()), ( - "teacher must stay frozen") + assert all(not p.requires_grad for p in teacher.transformer.parameters()), ("teacher must stay frozen") # Iteration 1 — first streaming chunk: 3 fresh latents, no overlap. method.optimizers_zero_grad(1) @@ -165,8 +154,7 @@ def test_streaming_long_tuning_multi_stage_train_steps( method.backward(loss_map, outputs, grad_accum_rounds=1) mask = outputs["dmd_latent_vis_dict"]["streaming_chunk_mask"] assert mask.dtype == torch.bool - assert int(mask.sum().item()) == 2, ( - "overlap latents must be masked out of the training chunk") + assert int(mask.sum().item()) == 2, ("overlap latents must be masked out of the training chunk") # Iteration 3 — final chunk fills the sequence (1 new latent). method.optimizers_zero_grad(3) @@ -255,8 +243,15 @@ class TestScheduleParsing: with pytest.raises(ValueError, match="ordered and non-overlapping"): parse_multi_phased_distill_schedule( [ - {"stage": "self_forcing", "start_step": 0, "end_step": 10}, - {"stage": "streaming_long", "start_step": 5}, + { + "stage": "self_forcing", + "start_step": 0, + "end_step": 10 + }, + { + "stage": "streaming_long", + "start_step": 5 + }, ], default_num_latent_t=6, default_streaming_chunk_size=3, @@ -293,12 +288,10 @@ class TestPadFirstFrameLatent: def test_passthrough_without_condition_or_when_long_enough(self) -> None: method = _method_without_init() raw: dict = {"other": 1} - assert method._with_padded_first_frame_latent( - raw, target_latent_frames=3) is raw + assert method._with_padded_first_frame_latent(raw, target_latent_frames=3) is raw full = {"first_frame_latent": torch.zeros(1, 4, 3, 2, 2)} - assert method._with_padded_first_frame_latent( - full, target_latent_frames=3) is full + assert method._with_padded_first_frame_latent(full, target_latent_frames=3) is full def test_rejects_non_5d_condition(self) -> None: method = _method_without_init() @@ -347,12 +340,10 @@ class TestStageChunkSize: method._chunk_size = 7 assert method._stage_chunk_size(_make_stage()) == 3 - assert method._stage_chunk_size( - _make_stage(streaming_chunk_size=None)) == 5 + assert method._stage_chunk_size(_make_stage(streaming_chunk_size=None)) == 5 method.method_config = {} - assert method._stage_chunk_size( - _make_stage(streaming_chunk_size=None)) == 7 + assert method._stage_chunk_size(_make_stage(streaming_chunk_size=None)) == 7 def test_rejects_non_positive_chunk_size(self) -> None: method = _method_without_init() diff --git a/fastvideo/tests/train/methods/test_wan_causal_cd.py b/fastvideo/tests/train/methods/test_wan_causal_cd.py index db514f7ed..3e4b3b47f 100644 --- a/fastvideo/tests/train/methods/test_wan_causal_cd.py +++ b/fastvideo/tests/train/methods/test_wan_causal_cd.py @@ -26,10 +26,7 @@ from fastvideo.train.methods.consistency_model.causal_cd import ( from fastvideo.train.models.wan import WanCausalModel from fastvideo.train.utils.config import load_run_config - -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" - / "wan_causal_t2v_causal_cd_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "wan_causal_t2v_causal_cd_min.yaml") def _build_synthetic_batch( @@ -38,18 +35,14 @@ def _build_synthetic_batch( ) -> dict[str, torch.Tensor]: batch_size = 1 return { - "text_embedding": - torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), - "text_attention_mask": - torch.ones(batch_size, 16, device=device), - "vae_latent": - torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype), + "text_embedding": torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), + "text_attention_mask": torch.ones(batch_size, 16, device=device), + "vae_latent": torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype), } @pytest.mark.usefixtures("distributed_setup") -def test_wan_causal_cd_single_train_step( - monkeypatch: pytest.MonkeyPatch) -> None: +def test_wan_causal_cd_single_train_step(monkeypatch: pytest.MonkeyPatch) -> None: if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -87,7 +80,11 @@ def test_wan_causal_cd_single_train_step( method = CausalConsistencyDistillationMethod( cfg=cfg, - role_models={"student": student, "teacher": teacher, "ema": ema}, + role_models={ + "student": student, + "teacher": teacher, + "ema": ema + }, ) method.on_train_start() @@ -96,8 +93,7 @@ def test_wan_causal_cd_single_train_step( loss = loss_map["total_loss"] assert torch.is_tensor(loss), "total_loss must be a torch.Tensor" - assert torch.isfinite(loss).item(), ( - f"total_loss is not finite: {loss.item()}") + assert torch.isfinite(loss).item(), (f"total_loss is not finite: {loss.item()}") method.backward(loss_map, outputs, grad_accum_rounds=1) @@ -108,15 +104,13 @@ def test_wan_causal_cd_single_train_step( assert len(trainable) > 0, "student layer 0 has no trainable parameters" for i, p in enumerate(trainable): assert p.grad is not None, f"student layer 0 param[{i}] has None grad" - assert torch.isfinite(p.grad).all().item(), ( - f"student layer 0 param[{i}] grad contains NaN/Inf") - assert any(p.grad.detach().float().norm().item() > 0.0 for p in trainable), ( - "all student layer-0 grads are exactly zero; consistency loss " - "did not reach the first transformer block") + assert torch.isfinite(p.grad).all().item(), (f"student layer 0 param[{i}] grad contains NaN/Inf") + assert any(p.grad.detach().float().norm().item() > 0.0 + for p in trainable), ("all student layer-0 grads are exactly zero; consistency loss " + "did not reach the first transformer block") # Teacher must stay frozen. - assert all(not p.requires_grad for p in teacher.transformer.parameters()), ( - "teacher must be frozen for Causal-CD") + assert all(not p.requires_grad for p in teacher.transformer.parameters()), ("teacher must be frozen for Causal-CD") # The EMA model and student start from the same checkpoint, so the first # parameter must match before any update. FSDP fully_shard params are @@ -127,8 +121,8 @@ def test_wan_causal_cd_single_train_step( ema_param = next(ema.transformer.parameters()) student_param = next(student.transformer.parameters()) - assert torch.equal(_local(ema_param), _local(student_param)), ( - "EMA model should start identical to the student (same checkpoint)") + assert torch.equal(_local(ema_param), + _local(student_param)), ("EMA model should start identical to the student (same checkpoint)") # The EMA update must move EMA toward the student. Apply a visibly large # perturbation so the bf16 lerp is well above rounding noise (the real @@ -138,10 +132,7 @@ def test_wan_causal_cd_single_train_step( before = _local(ema_param).detach().float().clone() method._update_ema() after = _local(ema_param).detach().float() - assert not torch.equal(before, after), ( - "EMA weights did not move after _update_ema") + assert not torch.equal(before, after), ("EMA weights did not move after _update_ema") # EMA = decay*ema + (1-decay)*student moves ~ (1-decay) of the gap. - expected = before + (1.0 - method._ema_decay) * ( - _local(student_param).detach().float() - before) - assert torch.allclose(after, expected, atol=1e-2), ( - "EMA update did not follow the expected lerp") + expected = before + (1.0 - method._ema_decay) * (_local(student_param).detach().float() - before) + assert torch.allclose(after, expected, atol=1e-2), ("EMA update did not follow the expected lerp") diff --git a/fastvideo/tests/train/methods/test_wan_causal_dfsft.py b/fastvideo/tests/train/methods/test_wan_causal_dfsft.py index b52164055..16dd606dd 100644 --- a/fastvideo/tests/train/methods/test_wan_causal_dfsft.py +++ b/fastvideo/tests/train/methods/test_wan_causal_dfsft.py @@ -30,10 +30,7 @@ from fastvideo.train.utils.config import load_run_config from .grad_norm_regression import check_grad_norm_regression - -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" - / "wan_causal_t2v_dfsft_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "wan_causal_t2v_dfsft_min.yaml") def _build_synthetic_batch( @@ -48,21 +45,17 @@ def _build_synthetic_batch( """ batch_size = 1 return { - "text_embedding": - torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), + "text_embedding": torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), # float32 mask, matching the production dataloader # (``fastvideo/dataset/utils.py`` emits ``torch.ones``/``zeros``). # ``prepare_batch`` casts to training dtype either way. - "text_attention_mask": - torch.ones(batch_size, 16, device=device), - "vae_latent": - torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype), + "text_attention_mask": torch.ones(batch_size, 16, device=device), + "vae_latent": torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype), } @pytest.mark.usefixtures("distributed_setup") -def test_wan_causal_dfsft_single_train_step( - monkeypatch: pytest.MonkeyPatch) -> None: +def test_wan_causal_dfsft_single_train_step(monkeypatch: pytest.MonkeyPatch) -> None: if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -101,14 +94,12 @@ def test_wan_causal_dfsft_single_train_step( loss = loss_map["total_loss"] assert torch.is_tensor(loss), "total_loss must be a torch.Tensor" - assert torch.isfinite(loss).item(), ( - f"total_loss is not finite: {loss.item()}") + assert torch.isfinite(loss).item(), (f"total_loss is not finite: {loss.item()}") method.backward(loss_map, outputs, grad_accum_rounds=1) blocks = getattr(model.transformer, "blocks", None) - assert blocks is not None and len(blocks) > 0, ( - "CausalWanTransformer is expected to expose ``.blocks``") + assert blocks is not None and len(blocks) > 0, ("CausalWanTransformer is expected to expose ``.blocks``") layer0 = blocks[0] trainable = [p for p in layer0.parameters() if p.requires_grad] @@ -116,14 +107,11 @@ def test_wan_causal_dfsft_single_train_step( for i, p in enumerate(trainable): assert p.grad is not None, f"layer 0 param[{i}] has None grad" - assert torch.isfinite(p.grad).all().item(), ( - f"layer 0 param[{i}] grad contains NaN/Inf") + assert torch.isfinite(p.grad).all().item(), (f"layer 0 param[{i}] grad contains NaN/Inf") - any_nonzero = any( - p.grad.detach().float().norm().item() > 0.0 for p in trainable) - assert any_nonzero, ( - "all layer-0 grads are exactly zero; backward did not " - "reach the first transformer block") + any_nonzero = any(p.grad.detach().float().norm().item() > 0.0 for p in trainable) + assert any_nonzero, ("all layer-0 grads are exactly zero; backward did not " + "reach the first transformer block") # 5a-ii: device-keyed grad-norm regression on top of the same harness. # Skips when the current GPU has no seeded reference. @@ -133,6 +121,4 @@ def test_wan_causal_dfsft_single_train_step( # observed 3.2562 vs 3.5860 (10.13% apart) with identical code, straddling # the default 10%. 12% covers both winners; real wiring breakage (dead # grads, scale bugs) still lands far outside it. - check_grad_norm_regression("test_wan_causal_dfsft", - model.transformer, - rtol=0.12) + check_grad_norm_regression("test_wan_causal_dfsft", model.transformer, rtol=0.12) diff --git a/fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py b/fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py index e91186a5f..77e21732c 100644 --- a/fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py +++ b/fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py @@ -25,10 +25,7 @@ from fastvideo.train.methods.fine_tuning.dfsft import ( from fastvideo.train.models.wan import WanCausalModel from fastvideo.train.utils.config import load_run_config - -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" - / "wan_causal_t2v_dfsft_framewise_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "wan_causal_t2v_dfsft_framewise_min.yaml") def _build_synthetic_batch( @@ -37,18 +34,14 @@ def _build_synthetic_batch( ) -> dict[str, torch.Tensor]: batch_size = 1 return { - "text_embedding": - torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), - "text_attention_mask": - torch.ones(batch_size, 16, device=device), - "vae_latent": - torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype), + "text_embedding": torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), + "text_attention_mask": torch.ones(batch_size, 16, device=device), + "vae_latent": torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype), } @pytest.mark.usefixtures("distributed_setup") -def test_wan_causal_dfsft_framewise_single_train_step( - monkeypatch: pytest.MonkeyPatch) -> None: +def test_wan_causal_dfsft_framewise_single_train_step(monkeypatch: pytest.MonkeyPatch) -> None: if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -70,8 +63,7 @@ def test_wan_causal_dfsft_framewise_single_train_step( trainable=True, num_frames_per_block=student_cfg.get("num_frames_per_block"), ) - assert model.transformer.num_frame_per_block == 1, ( - "frame-wise override did not reach the transformer") + assert model.transformer.num_frame_per_block == 1, ("frame-wise override did not reach the transformer") model.transformer = model.transformer.to(device=device, dtype=dtype) method = DiffusionForcingSFTMethod( @@ -85,8 +77,7 @@ def test_wan_causal_dfsft_framewise_single_train_step( loss = loss_map["total_loss"] assert torch.is_tensor(loss), "total_loss must be a torch.Tensor" - assert torch.isfinite(loss).item(), ( - f"total_loss is not finite: {loss.item()}") + assert torch.isfinite(loss).item(), (f"total_loss is not finite: {loss.item()}") method.backward(loss_map, outputs, grad_accum_rounds=1) @@ -99,11 +90,8 @@ def test_wan_causal_dfsft_framewise_single_train_step( for i, p in enumerate(trainable): assert p.grad is not None, f"layer 0 param[{i}] has None grad" - assert torch.isfinite(p.grad).all().item(), ( - f"layer 0 param[{i}] grad contains NaN/Inf") + assert torch.isfinite(p.grad).all().item(), (f"layer 0 param[{i}] grad contains NaN/Inf") - any_nonzero = any( - p.grad.detach().float().norm().item() > 0.0 for p in trainable) - assert any_nonzero, ( - "all layer-0 grads are exactly zero; backward did not " - "reach the first transformer block") + any_nonzero = any(p.grad.detach().float().norm().item() > 0.0 for p in trainable) + assert any_nonzero, ("all layer-0 grads are exactly zero; backward did not " + "reach the first transformer block") diff --git a/fastvideo/tests/train/methods/test_wan_causal_tfsft.py b/fastvideo/tests/train/methods/test_wan_causal_tfsft.py index 36c7d8e9d..590a95ece 100644 --- a/fastvideo/tests/train/methods/test_wan_causal_tfsft.py +++ b/fastvideo/tests/train/methods/test_wan_causal_tfsft.py @@ -25,10 +25,7 @@ from fastvideo.train.methods.fine_tuning.tfsft import ( from fastvideo.train.models.wan import WanCausalModel from fastvideo.train.utils.config import load_run_config - -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" - / "wan_causal_t2v_tfsft_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "wan_causal_t2v_tfsft_min.yaml") def _build_synthetic_batch( @@ -37,18 +34,14 @@ def _build_synthetic_batch( ) -> dict[str, torch.Tensor]: batch_size = 1 return { - "text_embedding": - torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), - "text_attention_mask": - torch.ones(batch_size, 16, device=device), - "vae_latent": - torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype), + "text_embedding": torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), + "text_attention_mask": torch.ones(batch_size, 16, device=device), + "vae_latent": torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype), } @pytest.mark.usefixtures("distributed_setup") -def test_wan_causal_tfsft_single_train_step( - monkeypatch: pytest.MonkeyPatch) -> None: +def test_wan_causal_tfsft_single_train_step(monkeypatch: pytest.MonkeyPatch) -> None: if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -81,14 +74,12 @@ def test_wan_causal_tfsft_single_train_step( loss = loss_map["total_loss"] assert torch.is_tensor(loss), "total_loss must be a torch.Tensor" - assert torch.isfinite(loss).item(), ( - f"total_loss is not finite: {loss.item()}") + assert torch.isfinite(loss).item(), (f"total_loss is not finite: {loss.item()}") method.backward(loss_map, outputs, grad_accum_rounds=1) blocks = getattr(model.transformer, "blocks", None) - assert blocks is not None and len(blocks) > 0, ( - "CausalWanTransformer is expected to expose ``.blocks``") + assert blocks is not None and len(blocks) > 0, ("CausalWanTransformer is expected to expose ``.blocks``") layer0 = blocks[0] trainable = [p for p in layer0.parameters() if p.requires_grad] @@ -96,18 +87,13 @@ def test_wan_causal_tfsft_single_train_step( for i, p in enumerate(trainable): assert p.grad is not None, f"layer 0 param[{i}] has None grad" - assert torch.isfinite(p.grad).all().item(), ( - f"layer 0 param[{i}] grad contains NaN/Inf") + assert torch.isfinite(p.grad).all().item(), (f"layer 0 param[{i}] grad contains NaN/Inf") - any_nonzero = any( - p.grad.detach().float().norm().item() > 0.0 for p in trainable) - assert any_nonzero, ( - "all layer-0 grads are exactly zero; backward did not " - "reach the first transformer block") + any_nonzero = any(p.grad.detach().float().norm().item() > 0.0 for p in trainable) + assert any_nonzero, ("all layer-0 grads are exactly zero; backward did not " + "reach the first transformer block") # Teacher forcing must build its own (concatenated) attention mask and # must not have constructed the diffusion-forcing mask. - assert model.transformer.teacher_forcing_block_mask is not None, ( - "teacher-forcing mask was not constructed") - assert model.transformer.block_mask is None, ( - "diffusion-forcing mask should not be built on the TF path") + assert model.transformer.teacher_forcing_block_mask is not None, ("teacher-forcing mask was not constructed") + assert model.transformer.block_mask is None, ("diffusion-forcing mask should not be built on the TF path") diff --git a/fastvideo/tests/train/methods/test_wan_finetune.py b/fastvideo/tests/train/methods/test_wan_finetune.py index f5b25414d..5867587ff 100644 --- a/fastvideo/tests/train/methods/test_wan_finetune.py +++ b/fastvideo/tests/train/methods/test_wan_finetune.py @@ -38,10 +38,7 @@ from fastvideo.train.utils.config import load_run_config from .grad_norm_regression import check_grad_norm_regression - -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" - / "wan_t2v_finetune_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "wan_t2v_finetune_min.yaml") def _build_synthetic_batch( @@ -56,23 +53,19 @@ def _build_synthetic_batch( """ batch_size = 1 return { - "text_embedding": - torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), + "text_embedding": torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), # float32 mask, matching the production dataloader # (``fastvideo/dataset/utils.py`` emits ``torch.ones``/``zeros``). # ``prepare_batch`` casts to training dtype either way. - "text_attention_mask": - torch.ones(batch_size, 16, device=device), + "text_attention_mask": torch.ones(batch_size, 16, device=device), # vae_latent in (B, C, T, H, W); ``prepare_batch`` will # truncate T to ``training.data.num_latent_t``. - "vae_latent": - torch.randn(batch_size, 16, 4, 8, 8, device=device, dtype=dtype), + "vae_latent": torch.randn(batch_size, 16, 4, 8, 8, device=device, dtype=dtype), } @pytest.mark.usefixtures("distributed_setup") -def test_wan_finetune_single_train_step( - monkeypatch: pytest.MonkeyPatch) -> None: +def test_wan_finetune_single_train_step(monkeypatch: pytest.MonkeyPatch) -> None: if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -115,14 +108,12 @@ def test_wan_finetune_single_train_step( loss = loss_map["total_loss"] assert torch.is_tensor(loss), "total_loss must be a torch.Tensor" - assert torch.isfinite(loss).item(), ( - f"total_loss is not finite: {loss.item()}") + assert torch.isfinite(loss).item(), (f"total_loss is not finite: {loss.item()}") method.backward(loss_map, outputs, grad_accum_rounds=1) blocks = getattr(model.transformer, "blocks", None) - assert blocks is not None and len(blocks) > 0, ( - "Wan transformer is expected to expose ``.blocks``") + assert blocks is not None and len(blocks) > 0, ("Wan transformer is expected to expose ``.blocks``") layer0 = blocks[0] trainable = [p for p in layer0.parameters() if p.requires_grad] @@ -130,17 +121,14 @@ def test_wan_finetune_single_train_step( for i, p in enumerate(trainable): assert p.grad is not None, f"layer 0 param[{i}] has None grad" - assert torch.isfinite(p.grad).all().item(), ( - f"layer 0 param[{i}] grad contains NaN/Inf") + assert torch.isfinite(p.grad).all().item(), (f"layer 0 param[{i}] grad contains NaN/Inf") # At least one layer-0 grad must have non-zero L2 norm so we # catch the case where backward ran but the first block was # detached from the loss (silent connectivity bug). - any_nonzero = any( - p.grad.detach().float().norm().item() > 0.0 for p in trainable) - assert any_nonzero, ( - "all layer-0 grads are exactly zero; backward did not " - "reach the first transformer block") + any_nonzero = any(p.grad.detach().float().norm().item() > 0.0 for p in trainable) + assert any_nonzero, ("all layer-0 grads are exactly zero; backward did not " + "reach the first transformer block") # 5a-ii: device-keyed grad-norm regression on top of the same harness. # Skips when the current GPU has no seeded reference. diff --git a/fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py b/fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py index 2681ce8d0..e3db14651 100644 --- a/fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py +++ b/fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py @@ -59,9 +59,7 @@ from fastvideo.attention.selector import _cached_get_attn_backend from fastvideo.forward_context import set_forward_context from fastvideo.platforms import AttentionBackendEnum -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" / - "kandinsky5_t2v_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "kandinsky5_t2v_min.yaml") @pytest.mark.usefixtures("distributed_setup") @@ -108,9 +106,14 @@ def test_kandinsky5_attn_qat_train_engages_and_backprops(monkeypatch): latent_h = grid_h * patch_size[1] latent_w = grid_w * patch_size[2] - latents = torch.randn( - 1, latent_t, latent_h, latent_w, in_visual_dim, device=device, dtype=dtype, - requires_grad=True) + latents = torch.randn(1, + latent_t, + latent_h, + latent_w, + in_visual_dim, + device=device, + dtype=dtype, + requires_grad=True) if bool(getattr(transformer, "visual_cond", False)): # See Kandinsky5Model._build_distill_input_kwargs / # Kandinsky5LatentPreparationStage: visual_cond=True checkpoints @@ -150,9 +153,8 @@ def test_kandinsky5_attn_qat_train_engages_and_backprops(monkeypatch): assert latents.grad is not None assert torch.isfinite(latents.grad).all().item(), "input grad contains NaN/Inf" assert attn.to_query.weight.grad is not None - assert torch.isfinite(attn.to_query.weight.grad.to_local() - if hasattr(attn.to_query.weight.grad, "to_local") else - attn.to_query.weight.grad).all().item(), "weight grad contains NaN/Inf" + assert torch.isfinite(attn.to_query.weight.grad.to_local() if hasattr(attn.to_query.weight.grad, "to_local") + else attn.to_query.weight.grad).all().item(), "weight grad contains NaN/Inf" finally: # monkeypatch restores FASTVIDEO_ATTENTION_BACKEND itself, but the # selector cache is a separate process-wide functools.cache keyed on diff --git a/fastvideo/tests/train/models/test_load_cosmos.py b/fastvideo/tests/train/models/test_load_cosmos.py index 1dca9ef49..444aa5caf 100644 --- a/fastvideo/tests/train/models/test_load_cosmos.py +++ b/fastvideo/tests/train/models/test_load_cosmos.py @@ -31,8 +31,7 @@ from fastvideo.forward_context import set_forward_context from fastvideo.train.models.cosmos import CosmosModel from fastvideo.train.utils.config import load_run_config -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" / "cosmos_t2w_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "cosmos_t2w_min.yaml") # Cosmos 2.5 Reason1 (Qwen2.5-VL) text embedding width. _COSMOS_TEXT_DIM = 100352 @@ -63,11 +62,7 @@ def test_cosmos_model_loads_and_forwards(): # Small spatial + few frames so this fits next to the 2B model. b, c, t, h, w = 1, 16, 4, 32, 32 hidden_states = torch.randn(b, c, t, h, w, device=device, dtype=dtype) - encoder_hidden_states = torch.randn(b, - 16, - _COSMOS_TEXT_DIM, - device=device, - dtype=dtype) + encoder_hidden_states = torch.randn(b, 16, _COSMOS_TEXT_DIM, device=device, dtype=dtype) timestep = torch.tensor([[500]], device=device, dtype=dtype) condition_mask = torch.zeros(b, 1, t, h, w, device=device, dtype=dtype) padding_mask = torch.zeros(1, 1, h, w, device=device, dtype=dtype) @@ -87,7 +82,6 @@ def test_cosmos_model_loads_and_forwards(): if isinstance(out, tuple): out = out[0] - assert out.shape == hidden_states.shape, ( - f"output shape {tuple(out.shape)} != input shape " - f"{tuple(hidden_states.shape)}") + assert out.shape == hidden_states.shape, (f"output shape {tuple(out.shape)} != input shape " + f"{tuple(hidden_states.shape)}") assert torch.isfinite(out).all().item(), "output contains NaN/Inf" diff --git a/fastvideo/tests/train/models/test_load_hunyuan.py b/fastvideo/tests/train/models/test_load_hunyuan.py index 348e2113b..90d71d1c3 100644 --- a/fastvideo/tests/train/models/test_load_hunyuan.py +++ b/fastvideo/tests/train/models/test_load_hunyuan.py @@ -25,9 +25,7 @@ from fastvideo.forward_context import set_forward_context from fastvideo.train.models.hunyuan import HunyuanModel from fastvideo.train.utils.config import load_run_config -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" / - "hunyuan_t2v_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "hunyuan_t2v_min.yaml") @pytest.mark.usefixtures("distributed_setup") @@ -71,7 +69,6 @@ def test_hunyuan_model_loads_and_forwards(): if isinstance(out, tuple): out = out[0] - assert out.shape == hidden_states.shape, ( - f"output shape {tuple(out.shape)} != input shape " - f"{tuple(hidden_states.shape)}") + assert out.shape == hidden_states.shape, (f"output shape {tuple(out.shape)} != input shape " + f"{tuple(hidden_states.shape)}") assert torch.isfinite(out).all().item(), "output contains NaN/Inf" diff --git a/fastvideo/tests/train/models/test_load_kandinsky5.py b/fastvideo/tests/train/models/test_load_kandinsky5.py index 7d9ed6b98..2eff39833 100644 --- a/fastvideo/tests/train/models/test_load_kandinsky5.py +++ b/fastvideo/tests/train/models/test_load_kandinsky5.py @@ -28,9 +28,7 @@ from fastvideo.forward_context import set_forward_context from fastvideo.train.models.kandinsky5 import Kandinsky5Model from fastvideo.train.utils.config import load_run_config -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" / - "kandinsky5_t2v_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "kandinsky5_t2v_min.yaml") @pytest.mark.usefixtures("distributed_setup") @@ -68,8 +66,7 @@ def test_kandinsky5_model_loads_and_forwards(): latent_h = grid_h * patch_size[1] latent_w = grid_w * patch_size[2] - latents = torch.randn( - 1, latent_t, latent_h, latent_w, in_visual_dim, device=device, dtype=dtype) + latents = torch.randn(1, latent_t, latent_h, latent_w, in_visual_dim, device=device, dtype=dtype) if bool(getattr(transformer, "visual_cond", False)): # The shipped checkpoint has visual_cond=True (unified T2V/I2V): # Kandinsky5VisualEmbeddings expects [real | zero_cond | zero_mask] @@ -107,6 +104,5 @@ def test_kandinsky5_model_loads_and_forwards(): ).sample expected_shape = (1, latent_t, latent_h, latent_w, out_visual_dim) - assert tuple(out.shape) == expected_shape, ( - f"output shape {tuple(out.shape)} != expected {expected_shape}") + assert tuple(out.shape) == expected_shape, (f"output shape {tuple(out.shape)} != expected {expected_shape}") assert torch.isfinite(out).all().item(), "output contains NaN/Inf" diff --git a/fastvideo/tests/train/models/test_load_longcat.py b/fastvideo/tests/train/models/test_load_longcat.py index 778f0db99..f0fcdffd9 100644 --- a/fastvideo/tests/train/models/test_load_longcat.py +++ b/fastvideo/tests/train/models/test_load_longcat.py @@ -32,8 +32,7 @@ from fastvideo.forward_context import set_forward_context from fastvideo.train.models.longcat import LongCatModel from fastvideo.train.utils.config import load_run_config -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" / "longcat_t2v_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "longcat_t2v_min.yaml") # LongCat caption_channels (text embedding width). _LONGCAT_TEXT_DIM = 4096 @@ -62,11 +61,7 @@ def test_longcat_model_loads_and_forwards(): # LongCat transformer takes [B, C, T, H, W] (in_channels=16, # patch_size=(1,2,2)) and a [B, N_text, 4096] text embedding. hidden_states = torch.randn(1, 16, 4, 32, 32, device=device, dtype=dtype) - encoder_hidden_states = torch.randn(1, - 16, - _LONGCAT_TEXT_DIM, - device=device, - dtype=dtype) + encoder_hidden_states = torch.randn(1, 16, _LONGCAT_TEXT_DIM, device=device, dtype=dtype) encoder_attention_mask = torch.ones(1, 16, device=device, dtype=dtype) timestep = torch.tensor([500], device=device, dtype=dtype) @@ -83,7 +78,6 @@ def test_longcat_model_loads_and_forwards(): if isinstance(out, tuple): out = out[0] - assert out.shape == hidden_states.shape, ( - f"output shape {tuple(out.shape)} != input shape " - f"{tuple(hidden_states.shape)}") + assert out.shape == hidden_states.shape, (f"output shape {tuple(out.shape)} != input shape " + f"{tuple(hidden_states.shape)}") assert torch.isfinite(out).all().item(), "output contains NaN/Inf" diff --git a/fastvideo/tests/train/models/test_load_ltx2.py b/fastvideo/tests/train/models/test_load_ltx2.py index 05665bc4d..5fa4c02c6 100644 --- a/fastvideo/tests/train/models/test_load_ltx2.py +++ b/fastvideo/tests/train/models/test_load_ltx2.py @@ -61,13 +61,15 @@ class _TinyLTX2Transformer(torch.nn.Module): self.a2v_weight = torch.nn.Parameter(torch.tensor(1.0)) self.v2a_weight = torch.nn.Parameter(torch.tensor(1.0)) self.av_ca_weight = torch.nn.Parameter(torch.tensor(1.0)) - self.config = type("Config", (), { - "arch_config": type("Arch", (), { - "caption_proj_before_connector": is_ltx2_3, - "cross_attention_dim": 4096, - "caption_channels": 3840, - })(), - })() + self.config = type( + "Config", (), { + "arch_config": + type("Arch", (), { + "caption_proj_before_connector": is_ltx2_3, + "cross_attention_dim": 4096, + "caption_channels": 3840, + })(), + })() self.forward_fps: float | None = None self.forward_timestep: torch.Tensor | None = None self.forward_text_dim: int | None = None @@ -94,8 +96,7 @@ def _gpu_too_small() -> bool: return total < _MIN_GPU_MEMORY_GB * 1024**3 -def test_ltx2_rejects_audio_training_before_loading( - monkeypatch: pytest.MonkeyPatch) -> None: +def test_ltx2_rejects_audio_training_before_loading(monkeypatch: pytest.MonkeyPatch) -> None: cfg = load_run_config(str(_FIXTURE_DIR / _CASES["ltx2"][0])) def fail_if_loaded(*_args, **_kwargs): @@ -110,8 +111,7 @@ def test_ltx2_rejects_audio_training_before_loading( ) -def test_ltx2_forwards_role_attention_backend_to_loader( - monkeypatch: pytest.MonkeyPatch) -> None: +def test_ltx2_forwards_role_attention_backend_to_loader(monkeypatch: pytest.MonkeyPatch) -> None: cfg = load_run_config(str(_FIXTURE_DIR / _CASES["ltx2"][0])) transformer = _TinyLTX2Transformer(is_ltx2_3=False) captured: list[AttentionBackendEnum | None] = [] @@ -136,8 +136,7 @@ def test_ltx2_forwards_role_attention_backend_to_loader( @pytest.mark.parametrize("case", _CASES.keys()) -def test_ltx2_wrapper_contract_runs_on_cpu( - case: str, monkeypatch: pytest.MonkeyPatch) -> None: +def test_ltx2_wrapper_contract_runs_on_cpu(case: str, monkeypatch: pytest.MonkeyPatch) -> None: fixture_name, text_dim = _CASES[case] cfg = load_run_config(str(_FIXTURE_DIR / fixture_name)) transformer = _TinyLTX2Transformer(is_ltx2_3=case == "ltx2_3") @@ -173,7 +172,9 @@ def test_ltx2_wrapper_contract_runs_on_cpu( "text_embedding": torch.randn(1, 4, text_dim), "text_attention_mask": torch.ones(1, 4), "vae_latent": raw_latents, - "info_list": [{"fps": 12.0}], + "info_list": [{ + "fps": 12.0 + }], }, generator=torch.Generator(device="cpu").manual_seed(0), ) @@ -216,8 +217,7 @@ def test_ltx2_wrapper_contract_runs_on_cpu( backward_fps: list[float] = [] transformer.video_weight.register_hook( - lambda grad: backward_fps.append( - get_forward_context().forward_batch.fps) or grad) + lambda grad: backward_fps.append(get_forward_context().forward_batch.fps) or grad) with set_forward_context( current_timestep=torch.tensor([-1.0]), attn_metadata=None, @@ -263,11 +263,7 @@ def test_ltx2_model_loads_and_forwards(case: str): b, c, t, h, w = 1, 128, 3, 8, 8 tokens = t * h * w hidden_states = torch.randn(b, c, t, h, w, device=device, dtype=dtype) - encoder_hidden_states = torch.randn(b, - _LTX2_TEXT_LEN, - text_dim, - device=device, - dtype=dtype) + encoder_hidden_states = torch.randn(b, _LTX2_TEXT_LEN, text_dim, device=device, dtype=dtype) timestep = torch.full((b, tokens), 0.5, device=device, dtype=torch.float32) with torch.no_grad(), torch.autocast(device.type, dtype=dtype), \ @@ -285,7 +281,6 @@ def test_ltx2_model_loads_and_forwards(case: str): ) assert torch.is_tensor(out), f"expected a tensor output, got {type(out)}" - assert out.shape == hidden_states.shape, ( - f"denoised output shape {tuple(out.shape)} != input latent shape " - f"{tuple(hidden_states.shape)}") + assert out.shape == hidden_states.shape, (f"denoised output shape {tuple(out.shape)} != input latent shape " + f"{tuple(hidden_states.shape)}") assert torch.isfinite(out).all().item(), "forward output contains NaN/Inf" diff --git a/fastvideo/tests/train/models/test_load_matrixgame2.py b/fastvideo/tests/train/models/test_load_matrixgame2.py index f8d173e91..ac182db84 100644 --- a/fastvideo/tests/train/models/test_load_matrixgame2.py +++ b/fastvideo/tests/train/models/test_load_matrixgame2.py @@ -33,8 +33,7 @@ from fastvideo.forward_context import set_forward_context from fastvideo.train.models.matrixgame2.matrixgame2 import MatrixGame2Model from fastvideo.train.utils.config import load_run_config -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" / "matrixgame2_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "matrixgame2_min.yaml") # Matrix-Game 2.0 CLIP image-embedding width. _MG2_IMAGE_DIM = 1280 @@ -65,11 +64,9 @@ def test_matrixgame2_model_loads_and_forwards(): # transformer expects (16 noise + image/mask cond channels). in_channels = transformer.patch_embedding.proj.in_channels - hidden_states = torch.randn(1, in_channels, 4, 32, 32, device=device, - dtype=dtype) + hidden_states = torch.randn(1, in_channels, 4, 32, 32, device=device, dtype=dtype) # CLIP image embeds drive the image-action cross-attention; no text. - encoder_hidden_states_image = torch.randn(1, 257, _MG2_IMAGE_DIM, - device=device, dtype=dtype) + encoder_hidden_states_image = torch.randn(1, 257, _MG2_IMAGE_DIM, device=device, dtype=dtype) timestep = torch.tensor([500], device=device, dtype=torch.long) with torch.no_grad(), set_forward_context( @@ -90,7 +87,6 @@ def test_matrixgame2_model_loads_and_forwards(): out = out[0] # out has out_channels (16); compare the non-channel dims. assert out.shape[0] == hidden_states.shape[0] - assert out.shape[2:] == hidden_states.shape[2:], ( - f"output spatial/temporal shape {tuple(out.shape)} != input " - f"{tuple(hidden_states.shape)}") + assert out.shape[2:] == hidden_states.shape[2:], (f"output spatial/temporal shape {tuple(out.shape)} != input " + f"{tuple(hidden_states.shape)}") assert torch.isfinite(out).all().item(), "output contains NaN/Inf" diff --git a/fastvideo/tests/train/models/test_load_wan.py b/fastvideo/tests/train/models/test_load_wan.py index 0099c1443..5b6c63a62 100644 --- a/fastvideo/tests/train/models/test_load_wan.py +++ b/fastvideo/tests/train/models/test_load_wan.py @@ -26,8 +26,7 @@ from fastvideo.forward_context import set_forward_context from fastvideo.train.models.wan import WanModel from fastvideo.train.utils.config import load_run_config -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" / "wan_t2v_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "wan_t2v_min.yaml") @pytest.mark.usefixtures("distributed_setup") @@ -54,11 +53,7 @@ def test_wan_model_loads_and_forwards(): # (Wan patch_size=(1,2,2)). Keeps activations small so the test # fits comfortably alongside the 1.3B model on a single L40S. hidden_states = torch.randn(1, 16, 4, 32, 32, device=device, dtype=dtype) - encoder_hidden_states = torch.randn(1, - 16, - 4096, - device=device, - dtype=dtype) + encoder_hidden_states = torch.randn(1, 16, 4096, device=device, dtype=dtype) encoder_attention_mask = torch.ones(1, 16, device=device, dtype=dtype) timestep = torch.tensor([500], device=device, dtype=dtype) @@ -76,7 +71,6 @@ def test_wan_model_loads_and_forwards(): if isinstance(out, tuple): out = out[0] - assert out.shape == hidden_states.shape, ( - f"output shape {tuple(out.shape)} != input shape " - f"{tuple(hidden_states.shape)}") + assert out.shape == hidden_states.shape, (f"output shape {tuple(out.shape)} != input shape " + f"{tuple(hidden_states.shape)}") assert torch.isfinite(out).all().item(), "output contains NaN/Inf" diff --git a/fastvideo/tests/train/models/test_load_wan_causal.py b/fastvideo/tests/train/models/test_load_wan_causal.py index 6afb07076..ec8321de1 100644 --- a/fastvideo/tests/train/models/test_load_wan_causal.py +++ b/fastvideo/tests/train/models/test_load_wan_causal.py @@ -27,9 +27,7 @@ import torch from fastvideo.train.models.wan import WanCausalModel from fastvideo.train.utils.config import load_run_config -_FIXTURE = str( - Path(__file__).resolve().parent.parent / "fixtures" / - "wan_causal_t2v_min.yaml") +_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "wan_causal_t2v_min.yaml") @pytest.mark.usefixtures("distributed_setup") @@ -51,6 +49,5 @@ def test_wan_causal_model_loads(): # Spot-check that the override pulled the causal class, not the # plain Wan one. Transformer is a torch.nn.Module wrapping the # CausalWanTransformer3DModel architecture. - assert "Causal" in type(transformer).__name__, ( - f"expected CausalWanTransformer3DModel-derived class, got " - f"{type(transformer).__name__}") + assert "Causal" in type(transformer).__name__, (f"expected CausalWanTransformer3DModel-derived class, got " + f"{type(transformer).__name__}") diff --git a/fastvideo/tests/train/trainer/test_validation.py b/fastvideo/tests/train/trainer/test_validation.py index c2e2f0d6a..5a08a36f8 100644 --- a/fastvideo/tests/train/trainer/test_validation.py +++ b/fastvideo/tests/train/trainer/test_validation.py @@ -84,9 +84,7 @@ class _DummyMethod: self.weight.grad = None -def test_trainer_runs_validation_callback_during_training( - monkeypatch, -) -> None: +def test_trainer_runs_validation_callback_during_training(monkeypatch, ) -> None: tracker = _DummyTracker() group = SimpleNamespace(rank=0, local_rank=0, rank_in_group=0, world_size=1) @@ -124,7 +122,9 @@ def test_trainer_runs_validation_callback_during_training( trainer.run( method, - dataloader=[{"sample": "x"}], + dataloader=[{ + "sample": "x" + }], max_steps=3, ) diff --git a/fastvideo/tests/train/utils/test_checkpoint.py b/fastvideo/tests/train/utils/test_checkpoint.py index 663856f15..1bbab529c 100644 --- a/fastvideo/tests/train/utils/test_checkpoint.py +++ b/fastvideo/tests/train/utils/test_checkpoint.py @@ -24,7 +24,6 @@ from fastvideo.train.utils.checkpoint import ( _resolve_resume_checkpoint, ) - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -163,8 +162,7 @@ def test_find_latest_skips_non_checkpoint_dirs(tmp_path: Path) -> None: # --------------------------------------------------------------------------- -def test_resolve_latest_with_no_checkpoints_returns_none( - tmp_path: Path) -> None: +def test_resolve_latest_with_no_checkpoints_returns_none(tmp_path: Path) -> None: out = tmp_path / "outputs" out.mkdir() assert _resolve_resume_checkpoint("latest", output_dir=str(out)) is None @@ -180,8 +178,7 @@ def test_resolve_latest_returns_latest_checkpoint(tmp_path: Path) -> None: def test_resolve_explicit_checkpoint_dir(tmp_path: Path) -> None: ckpt = _make_checkpoint_dir(tmp_path, 42) - resolved = _resolve_resume_checkpoint(str(ckpt), - output_dir=str(tmp_path)) + resolved = _resolve_resume_checkpoint(str(ckpt), output_dir=str(tmp_path)) assert resolved is not None assert resolved.name == "checkpoint-42" @@ -189,8 +186,7 @@ def test_resolve_explicit_checkpoint_dir(tmp_path: Path) -> None: def test_resolve_dcp_subdir_returns_parent_checkpoint(tmp_path: Path) -> None: ckpt = _make_checkpoint_dir(tmp_path, 42) dcp_path = ckpt / "dcp" - resolved = _resolve_resume_checkpoint(str(dcp_path), - output_dir=str(tmp_path)) + resolved = _resolve_resume_checkpoint(str(dcp_path), output_dir=str(tmp_path)) assert resolved is not None assert resolved.name == "checkpoint-42" @@ -207,8 +203,7 @@ def test_resolve_output_dir_returns_latest(tmp_path: Path) -> None: def test_resolve_nonexistent_path_raises(tmp_path: Path) -> None: with pytest.raises(FileNotFoundError): - _resolve_resume_checkpoint(str(tmp_path / "missing"), - output_dir=str(tmp_path)) + _resolve_resume_checkpoint(str(tmp_path / "missing"), output_dir=str(tmp_path)) def test_resolve_checkpoint_without_dcp_raises(tmp_path: Path) -> None: diff --git a/fastvideo/tests/train/utils/test_config.py b/fastvideo/tests/train/utils/test_config.py index 0025d09ff..8471eea8c 100644 --- a/fastvideo/tests/train/utils/test_config.py +++ b/fastvideo/tests/train/utils/test_config.py @@ -27,8 +27,7 @@ def _minimal_yaml() -> dict[str, Any]: }, }, "method": { - "_target_": - "fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod", + "_target_": "fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod", }, "training": {}, } @@ -44,8 +43,7 @@ def test_minimal_yaml_loads_happy_path(tmp_path: Path) -> None: assert isinstance(cfg, RunConfig) assert isinstance(cfg.training, TrainingConfig) - assert cfg.method["_target_"] == ( - "fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod") + assert cfg.method["_target_"] == ("fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod") assert "student" in cfg.models assert cfg.callbacks == {} # raw retains the original YAML dict for downstream logging. @@ -252,17 +250,11 @@ def test_pipeline_quant_config_resolves_for_load_and_export(tmp_path: Path) -> N _run_config_from_raw, ) data = _minimal_yaml() - data["models"]["student"]["init_from"] = ( - "FastVideo/LTX2-Distilled-Diffusers") - data["pipeline"] = { - "dit_config": { - "quant_config": "nvfp4_qat_train" - } - } + data["models"]["student"]["init_from"] = ("FastVideo/LTX2-Distilled-Diffusers") + data["pipeline"] = {"dit_config": {"quant_config": "nvfp4_qat_train"}} cfg = load_run_config(_write_yaml(tmp_path, data)) - assert isinstance(cfg.training.pipeline_config.dit_config.quant_config, - NVFP4QATTrainConfig) + assert isinstance(cfg.training.pipeline_config.dit_config.quant_config, NVFP4QATTrainConfig) export_cfg = _run_config_from_raw(cfg.raw) assert isinstance( @@ -273,13 +265,8 @@ def test_pipeline_quant_config_resolves_for_load_and_export(tmp_path: Path) -> N def test_pipeline_quant_config_rejects_unknown_name(tmp_path: Path) -> None: data = _minimal_yaml() - data["models"]["student"]["init_from"] = ( - "FastVideo/LTX2-Distilled-Diffusers") - data["pipeline"] = { - "dit_config": { - "quant_config": "not_a_quantization_method" - } - } + data["models"]["student"]["init_from"] = ("FastVideo/LTX2-Distilled-Diffusers") + data["pipeline"] = {"dit_config": {"quant_config": "not_a_quantization_method"}} with pytest.raises(ValueError, match="Invalid quantization method"): load_run_config(_write_yaml(tmp_path, data)) @@ -387,8 +374,7 @@ def test_callbacks_passed_through_when_present(tmp_path: Path) -> None: data = _minimal_yaml() data["callbacks"] = { "grad_clip": { - "_target_": - "fastvideo.train.callbacks.grad_clip.GradNormClipCallback", + "_target_": "fastvideo.train.callbacks.grad_clip.GradNormClipCallback", "max_grad_norm": 1.0, }, } diff --git a/fastvideo/tests/train/utils/test_finetune_config_construction.py b/fastvideo/tests/train/utils/test_finetune_config_construction.py index 1b01ea672..baaa5eed3 100644 --- a/fastvideo/tests/train/utils/test_finetune_config_construction.py +++ b/fastvideo/tests/train/utils/test_finetune_config_construction.py @@ -29,14 +29,13 @@ def _recipe_id(path: Path) -> str: def test_fine_tuning_recipe_constructs(recipe: Path) -> None: cfg = load_run_config(str(recipe)) - assert cfg.training.pipeline_config is not None, ( - f"{_recipe_id(recipe)} did not resolve a pipeline config; " - "check its `pipeline:` key and init_from registry match") + assert cfg.training.pipeline_config is not None, (f"{_recipe_id(recipe)} did not resolve a pipeline config; " + "check its `pipeline:` key and init_from registry match") target = cfg.models["student"]["_target_"] - assert isinstance(target, str) and target and "." in target, ( - f"{_recipe_id(recipe)} student _target_ is not importable-shaped: " - f"{target!r}") + assert isinstance( + target, str) and target and "." in target, (f"{_recipe_id(recipe)} student _target_ is not importable-shaped: " + f"{target!r}") def test_recipe_glob_found_recipes() -> None: diff --git a/fastvideo/tests/train/utils/test_validation_media.py b/fastvideo/tests/train/utils/test_validation_media.py index 0750700fd..161adf8d1 100644 --- a/fastvideo/tests/train/utils/test_validation_media.py +++ b/fastvideo/tests/train/utils/test_validation_media.py @@ -13,10 +13,7 @@ from fastvideo.train.utils.validation_media import write_validation_mp4 def _rgb_frames(frame_count: int) -> list[np.ndarray]: """Create distinct even-sized RGB frames for codec tests.""" - return [ - np.full((16, 24, 3), fill_value=frame_index * 10, dtype=np.uint8) - for frame_index in range(frame_count) - ] + return [np.full((16, 24, 3), fill_value=frame_index * 10, dtype=np.uint8) for frame_index in range(frame_count)] def _stereo_waveform(sample_count: int, sample_rate: int) -> np.ndarray: diff --git a/fastvideo/tests/training/VSA/test_training_loss_VSA.py b/fastvideo/tests/training/VSA/test_training_loss_VSA.py index 99d1d4b83..7f81a6d6a 100644 --- a/fastvideo/tests/training/VSA/test_training_loss_VSA.py +++ b/fastvideo/tests/training/VSA/test_training_loss_VSA.py @@ -23,87 +23,56 @@ reference_wandb_summary_file = "fastvideo/tests/training/VSA/h100_reference_wand NUM_NODES = "1" NUM_GPUS_PER_NODE = "2" + def run_worker(): """Worker function that will be run on each GPU""" # Create and populate args parser = FlexibleArgumentParser() parser = TrainingArgs.add_cli_args(parser) parser = FastVideoArgs.add_cli_args(parser) - + # Set the arguments as they are in finetune_v1_test.sh args = parser.parse_args([ - "--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--inference_mode", "False", - "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--data_path", "data/mini_dataset_i2v_VSA/combined_parquet_dataset", - "--validation_dataset_file", "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json", - "--train_batch_size", "1", - "--num_latent_t", "4", - "--num_gpus", "2", - "--sp_size", "2", - "--tp_size", "2", - "--hsdp_replicate_dim", "1", - "--hsdp_shard_dim", "2", - "--train_sp_batch_size", "1", - "--dataloader_num_workers", "4", - "--gradient_accumulation_steps", "2", - "--max_train_steps", "5", - "--learning_rate", "1e-5", - "--mixed_precision", "bf16", - "--weight_only_checkpointing_steps", "30", - "--training_state_checkpointing_steps", "30", - "--validation_steps", "10", - "--validation_sampling_steps", "50", - "--log_validation", - "--checkpoints_total_limit", "3", - "--ema_start_step", "0", - "--training_cfg_rate", "0.0", - "--output_dir", "data/wan_finetune_test_VSA", - "--tracker_project_name", "wan_finetune_ci_VSA", - "--wandb_run_name", wandb_name, - "--num_height", "384", - "--num_width", "512", - "--num_frames", "13", - "--flow_shift", "3", - "--validation_guidance_scale", "3.0", - "--num_euler_timesteps", "50", - "--weight_decay", "0.01", - "--dit_precision", "fp32", - "--max_grad_norm", "1.0", - "--VSA_decay_rate", "0.01", - "--VSA_decay_interval_steps", "1", - "--VSA_sparsity", "0.9" + "--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--inference_mode", "False", + "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--data_path", + "data/mini_dataset_i2v_VSA/combined_parquet_dataset", "--validation_dataset_file", + "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json", "--train_batch_size", "1", + "--num_latent_t", "4", "--num_gpus", "2", "--sp_size", "2", "--tp_size", "2", "--hsdp_replicate_dim", "1", + "--hsdp_shard_dim", "2", "--train_sp_batch_size", "1", "--dataloader_num_workers", "4", + "--gradient_accumulation_steps", "2", "--max_train_steps", "5", "--learning_rate", "1e-5", "--mixed_precision", + "bf16", "--weight_only_checkpointing_steps", "30", "--training_state_checkpointing_steps", "30", + "--validation_steps", "10", "--validation_sampling_steps", "50", "--log_validation", + "--checkpoints_total_limit", "3", "--ema_start_step", "0", "--training_cfg_rate", "0.0", "--output_dir", + "data/wan_finetune_test_VSA", "--tracker_project_name", "wan_finetune_ci_VSA", "--wandb_run_name", wandb_name, + "--num_height", "384", "--num_width", "512", "--num_frames", "13", "--flow_shift", "3", + "--validation_guidance_scale", "3.0", "--num_euler_timesteps", "50", "--weight_decay", "0.01", + "--dit_precision", "fp32", "--max_grad_norm", "1.0", "--VSA_decay_rate", "0.01", "--VSA_decay_interval_steps", + "1", "--VSA_sparsity", "0.9" ]) - + # Call the main training function main(args) + def test_distributed_training(): """Test the distributed training setup""" os.environ["WANDB_MODE"] = "online" data_dir = Path("data/mini_dataset_i2v_VSA") - + if not data_dir.exists(): print(f"Downloading test dataset to {data_dir}...") - snapshot_download( - repo_id="BrianChen1129/mini_dataset_i2v_VSA", - local_dir=str(data_dir), - repo_type="dataset", - local_dir_use_symlinks=False - ) + snapshot_download(repo_id="BrianChen1129/mini_dataset_i2v_VSA", + local_dir=str(data_dir), + repo_type="dataset", + local_dir_use_symlinks=False) # Get the current file path current_file = Path(__file__).resolve() - + # Run torchrun command - cmd = [ - "torchrun", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE, - str(current_file) - ] - + cmd = ["torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, str(current_file)] + process = subprocess.run(cmd, check=True) summary_file = 'data/wan_finetune_test_VSA/tracker/wandb/latest-run/files/wandb-summary.json' @@ -116,12 +85,7 @@ def test_distributed_training(): reference_wandb_summary = json.load(open(reference_wandb_summary_file)) wandb_summary = json.load(open(summary_file)) - fields_and_thresholds = { - 'avg_step_time': 15, - 'grad_norm': 0.5, - 'step_time': 15, - 'train_loss': 0.04 - } + fields_and_thresholds = {'avg_step_time': 15, 'grad_norm': 0.5, 'step_time': 15, 'train_loss': 0.04} failures = [] for field, threshold in fields_and_thresholds.items(): @@ -130,11 +94,14 @@ def test_distributed_training(): diff = abs(ref_value - current_value) print(f"INFO: {field}, diff: {diff}, threshold: {threshold}, reference: {ref_value}, current: {current_value}") if diff > threshold: - failures.append(f"FAILED: {field} difference {diff} exceeds threshold of {threshold} (reference: {ref_value}, current: {current_value})") + failures.append( + f"FAILED: {field} difference {diff} exceeds threshold of {threshold} (reference: {ref_value}, current: {current_value})" + ) if failures: raise AssertionError("\n".join(failures)) + if __name__ == "__main__": if os.environ.get("LOCAL_RANK") is not None: # We're being run by torchrun diff --git a/fastvideo/tests/training/Vanilla/mfu_calculation.py b/fastvideo/tests/training/Vanilla/mfu_calculation.py index 2e4596b5e..33f3d76f4 100644 --- a/fastvideo/tests/training/Vanilla/mfu_calculation.py +++ b/fastvideo/tests/training/Vanilla/mfu_calculation.py @@ -41,91 +41,123 @@ def run_worker(): parser = FlexibleArgumentParser() parser = TrainingArgs.add_cli_args(parser) parser = FastVideoArgs.add_cli_args(parser) - + # Set the arguments as they are in finetune_t2v.sh args = parser.parse_args([ - "--model_path", MODEL_PATH, - "--inference_mode", "False", - "--pretrained_model_name_or_path", MODEL_PATH, - "--data_path", DATA_PATH, - "--dataloader_num_workers", "1", - "--train_batch_size", "4", - "--train_sp_batch_size", "1", - "--gradient_accumulation_steps", GRAD_ACCUM, - "--num_latent_t", "20", - "--num_height", "720", - "--num_width", "1280", - "--num_frames", "77", - "--enable_gradient_checkpointing_type", "full", - "--max_train_steps", "20", - "--learning_rate", "5e-5", - "--mixed_precision", "bf16", - "--weight_only_checkpointing_steps", "250", - "--training_state_checkpointing_steps", "250", - "--weight_decay", "1e-4", - "--max_grad_norm", "1.0", - "--num_euler_timesteps", "50", - "--multi_phased_distill_schedule", "4000-1", + "--model_path", + MODEL_PATH, + "--inference_mode", + "False", + "--pretrained_model_name_or_path", + MODEL_PATH, + "--data_path", + DATA_PATH, + "--dataloader_num_workers", + "1", + "--train_batch_size", + "4", + "--train_sp_batch_size", + "1", + "--gradient_accumulation_steps", + GRAD_ACCUM, + "--num_latent_t", + "20", + "--num_height", + "720", + "--num_width", + "1280", + "--num_frames", + "77", + "--enable_gradient_checkpointing_type", + "full", + "--max_train_steps", + "20", + "--learning_rate", + "5e-5", + "--mixed_precision", + "bf16", + "--weight_only_checkpointing_steps", + "250", + "--training_state_checkpointing_steps", + "250", + "--weight_decay", + "1e-4", + "--max_grad_norm", + "1.0", + "--num_euler_timesteps", + "50", + "--multi_phased_distill_schedule", + "4000-1", "--not_apply_cfg_solver", - "--training_cfg_rate", "0.1", - "--ema_start_step", "0", - "--dit_precision", "fp32", - "--output_dir", str(OUTPUT_DIR), - "--tracker_project_name", "wan_t2v_finetune", - "--checkpoints_total_limit", "3", - "--validation_dataset_file", VALIDATION_DATASET_FILE, - "--validation_steps", "200", - "--validation_sampling_steps", "50", - "--validation_guidance_scale", "6.0", + "--training_cfg_rate", + "0.1", + "--ema_start_step", + "0", + "--dit_precision", + "fp32", + "--output_dir", + str(OUTPUT_DIR), + "--tracker_project_name", + "wan_t2v_finetune", + "--checkpoints_total_limit", + "3", + "--validation_dataset_file", + VALIDATION_DATASET_FILE, + "--validation_steps", + "200", + "--validation_sampling_steps", + "50", + "--validation_guidance_scale", + "6.0", #"--enable_torch_compile", #"--log_validation", - "--num_gpus", NUM_GPUS_PER_NODE, - "--sp_size", NUM_GPUS_PER_NODE, - "--tp_size", "1", - "--hsdp_replicate_dim", NUM_GPUS_PER_NODE, - "--hsdp_shard_dim", "1" + "--num_gpus", + NUM_GPUS_PER_NODE, + "--sp_size", + NUM_GPUS_PER_NODE, + "--tp_size", + "1", + "--hsdp_replicate_dim", + NUM_GPUS_PER_NODE, + "--hsdp_shard_dim", + "1" ]) # Call the main training function - pipeline = WanTrainingPipeline.from_pretrained( - args.pretrained_model_name_or_path, args=args) + pipeline = WanTrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args) args = pipeline.training_args pipeline.train() logger.info("Training pipeline done") + def test_distributed_training(): """Test the distributed training setup""" os.environ["WANDB_MODE"] = "online" data_dir = Path("data/crush-smol_processed_t2v") - + if not data_dir.exists(): print(f"Downloading test dataset to {data_dir}...") - snapshot_download( - repo_id="wlsaidhi/crush-smol_processed_t2v", - local_dir=str(data_dir), - repo_type="dataset", - local_dir_use_symlinks=False - ) - + snapshot_download(repo_id="wlsaidhi/crush-smol_processed_t2v", + local_dir=str(data_dir), + repo_type="dataset", + local_dir_use_symlinks=False) + # Get the current file path current_file = Path(__file__).resolve() - + # Run torchrun command cmd = [ - "torchrun", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE, - "--master_port", MASTER_PORT, + "torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, "--master_port", MASTER_PORT, str(current_file) ] process = subprocess.run(cmd, capture_output=True, text=True) - + # Print stdout and stderr for debugging if process.stdout: print("STDOUT:", process.stdout) if process.stderr: print("STDERR:", process.stderr) - + # Check if the process failed if process.returncode != 0: print(f"Process failed with return code: {process.returncode}") @@ -135,7 +167,7 @@ def test_distributed_training(): with summary_file.open() as f: wandb_summary = json.load(f) - + # Calculate and print MFU metrics device_name = torch.cuda.get_device_name() try: @@ -148,9 +180,6 @@ def test_distributed_training(): num_layers = wandb_summary.get("num_layers") ffn_dim = wandb_summary.get("ffn_dim") - - - # FLOPs per layer (forward pass) # - QKV + out proj: 8 * hidden_dim^2 * seq_len # - Cross-attn proj: 4 * hidden_dim^2 * seq_len + 4 * hidden_dim^2 * context_len @@ -158,24 +187,18 @@ def test_distributed_training(): # - Self-attn matmuls: 4 * seq_len^2 * hidden_dim # - Cross-attn matmuls: 4 * seq_len * context_len * hidden_dim qkv_out_flops = 8 * hidden_dim * hidden_dim * seq_len - cross_attn_proj_flops = ( - (4 * hidden_dim * hidden_dim * seq_len) + - (4 * hidden_dim * hidden_dim * context_len) - ) + cross_attn_proj_flops = ((4 * hidden_dim * hidden_dim * seq_len) + (4 * hidden_dim * hidden_dim * context_len)) mlp_flops = 4 * hidden_dim * ffn_dim * seq_len self_attn_flops = 4 * seq_len * seq_len * hidden_dim cross_attn_flops = 4 * seq_len * context_len * hidden_dim - flops_per_layer = ( - qkv_out_flops + cross_attn_proj_flops + mlp_flops + self_attn_flops + cross_attn_flops - ) + flops_per_layer = (qkv_out_flops + cross_attn_proj_flops + mlp_flops + self_attn_flops + cross_attn_flops) # With full activation checkpointing: 1 forward + 3 backward (1 recompute + 2 gradient) achieved_flops = batch_size * flops_per_layer * num_layers * 4 - # Account for gradient accumulation (from config) grad_accum = int(GRAD_ACCUM) - achieved_flops *= grad_accum + achieved_flops *= grad_accum # Peak FLOPs based on device if "H100" in device_name: @@ -188,11 +211,11 @@ def test_distributed_training(): peak_flops_per_gpu = 362e12 else: raise ValueError(f"Device {device_name} not supported") - + # Total peak (2 GPUs) world_size = int(NUM_GPUS_PER_NODE) total_peak_flops = peak_flops_per_gpu * world_size - + # Calculate MFU achieved_flops_per_sec = achieved_flops / avg_step_time if avg_step_time > 0 else 0 mfu = (achieved_flops_per_sec / total_peak_flops * 100) if total_peak_flops > 0 else 0 @@ -200,7 +223,7 @@ def test_distributed_training(): print(f"Per-Step MFU: {mfu:.4f}%") except Exception as e: print(f"Could not calculate MFU: {e}") - + if __name__ == "__main__": if os.environ.get("LOCAL_RANK") is not None: diff --git a/fastvideo/tests/training/Vanilla/test_training_loss.py b/fastvideo/tests/training/Vanilla/test_training_loss.py index 4e1fa5daf..c271b69e8 100644 --- a/fastvideo/tests/training/Vanilla/test_training_loss.py +++ b/fastvideo/tests/training/Vanilla/test_training_loss.py @@ -1,4 +1,5 @@ import os + os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29512" import sys @@ -30,91 +31,61 @@ def run_worker(): parser = FlexibleArgumentParser() parser = TrainingArgs.add_cli_args(parser) parser = FastVideoArgs.add_cli_args(parser) - + # Set the arguments as they are in finetune_v1_test.sh args = parser.parse_args([ - "--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--inference_mode", "False", - "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--data_path", "data/crush-smol_processed_t2v/combined_parquet_dataset", - "--validation_dataset_file", "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json", - "--train_batch_size", "2", - "--num_latent_t", "4", - "--num_gpus", "2", - "--sp_size", "2", - "--tp_size", "2", - "--hsdp_replicate_dim", "1", - "--hsdp_shard_dim", "2", - "--train_sp_batch_size", "1", - "--dataloader_num_workers", "1", - "--gradient_accumulation_steps", "2", - "--max_train_steps", "5", - "--learning_rate", "1e-6", - "--mixed_precision", "bf16", - "--weight_only_checkpointing_steps", "30", - "--training_state_checkpointing_steps", "30", - "--validation_steps", "10", - "--validation_sampling_steps", "8", - "--log_validation", - "--checkpoints_total_limit", "3", - "--ema_start_step", "0", - "--training_cfg_rate", "0.0", - "--output_dir", "data/wan_finetune_test", - "--tracker_project_name", "wan_finetune_ci", - "--wandb_run_name", wandb_name, - "--num_height", "480", - "--num_width", "832", - "--num_frames", "81", - "--flow_shift", "3", - "--validation_guidance_scale", "3.0", - "--num_euler_timesteps", "50", - "--multi_phased_distill_schedule", "4000-1", - "--weight_decay", "0.01", - "--not_apply_cfg_solver", - "--dit_precision", "fp32", - "--max_grad_norm", "1.0" + "--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--inference_mode", "False", + "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--data_path", + "data/crush-smol_processed_t2v/combined_parquet_dataset", "--validation_dataset_file", + "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json", "--train_batch_size", "2", + "--num_latent_t", "4", "--num_gpus", "2", "--sp_size", "2", "--tp_size", "2", "--hsdp_replicate_dim", "1", + "--hsdp_shard_dim", "2", "--train_sp_batch_size", "1", "--dataloader_num_workers", "1", + "--gradient_accumulation_steps", "2", "--max_train_steps", "5", "--learning_rate", "1e-6", "--mixed_precision", + "bf16", "--weight_only_checkpointing_steps", "30", "--training_state_checkpointing_steps", "30", + "--validation_steps", "10", "--validation_sampling_steps", "8", "--log_validation", "--checkpoints_total_limit", + "3", "--ema_start_step", "0", "--training_cfg_rate", "0.0", "--output_dir", "data/wan_finetune_test", + "--tracker_project_name", "wan_finetune_ci", "--wandb_run_name", wandb_name, "--num_height", "480", + "--num_width", "832", "--num_frames", "81", "--flow_shift", "3", "--validation_guidance_scale", "3.0", + "--num_euler_timesteps", "50", "--multi_phased_distill_schedule", "4000-1", "--weight_decay", "0.01", + "--not_apply_cfg_solver", "--dit_precision", "fp32", "--max_grad_norm", "1.0" ]) # Call the main training function - pipeline = WanTrainingPipeline.from_pretrained( - args.pretrained_model_name_or_path, args=args) + pipeline = WanTrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args) args = pipeline.training_args pipeline.train() logger.info("Training pipeline done") + def test_distributed_training(): """Test the distributed training setup""" os.environ["WANDB_MODE"] = "online" data_dir = Path("data/crush-smol_processed_t2v") - + if not data_dir.exists(): print(f"Downloading test dataset to {data_dir}...") - snapshot_download( - repo_id="wlsaidhi/crush-smol_processed_t2v", - local_dir=str(data_dir), - repo_type="dataset", - local_dir_use_symlinks=False - ) - + snapshot_download(repo_id="wlsaidhi/crush-smol_processed_t2v", + local_dir=str(data_dir), + repo_type="dataset", + local_dir_use_symlinks=False) + # Get the current file path current_file = Path(__file__).resolve() - + # Run torchrun command cmd = [ - "torchrun", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE, - "--master_port", os.environ["MASTER_PORT"], + "torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, "--master_port", + os.environ["MASTER_PORT"], str(current_file) ] process = subprocess.run(cmd, capture_output=True, text=True) - + # Print stdout and stderr for debugging if process.stdout: print("STDOUT:", process.stdout) if process.stderr: print("STDERR:", process.stderr) - + # Check if the process failed if process.returncode != 0: print(f"Process failed with return code: {process.returncode}") @@ -135,12 +106,7 @@ def test_distributed_training(): reference_wandb_summary = json.load(open(reference_wandb_summary_file)) wandb_summary = json.load(open(summary_file)) - fields_and_thresholds = { - 'avg_step_time': 15.0, - 'grad_norm': 0.3, - 'step_time': 15.0, - 'train_loss': 0.0025 - } + fields_and_thresholds = {'avg_step_time': 15.0, 'grad_norm': 0.3, 'step_time': 15.0, 'train_loss': 0.0025} failures = [] for field, threshold in fields_and_thresholds.items(): @@ -149,11 +115,14 @@ def test_distributed_training(): diff = abs(ref_value - current_value) print(f"INFO: {field}, diff: {diff}, threshold: {threshold}, reference: {ref_value}, current: {current_value}") if diff > threshold: - failures.append(f"FAILED: {field} difference {diff} exceeds threshold of {threshold} (reference: {ref_value}, current: {current_value})") + failures.append( + f"FAILED: {field} difference {diff} exceeds threshold of {threshold} (reference: {ref_value}, current: {current_value})" + ) if failures: raise AssertionError("\n".join(failures)) + if __name__ == "__main__": if os.environ.get("LOCAL_RANK") is not None: # We're being run by torchrun diff --git a/fastvideo/tests/training/checkpoint/test_checkpointing_utils.py b/fastvideo/tests/training/checkpoint/test_checkpointing_utils.py index 086da6732..eb3649bc6 100644 --- a/fastvideo/tests/training/checkpoint/test_checkpointing_utils.py +++ b/fastvideo/tests/training/checkpoint/test_checkpointing_utils.py @@ -5,6 +5,7 @@ from fastvideo.training.checkpointing_utils import ModelWrapper class DummyWrappedModule(nn.Module): + def __init__(self): super().__init__() self._lora_a = nn.Parameter(torch.tensor([1.0, 2.0])) @@ -42,4 +43,4 @@ def test_model_wrapper_filters_wrapped_trainable_params(monkeypatch): assert set(filtered_state_dict.keys()) == {"layer.lora_A", "layer.lora_B"} assert torch.equal(filtered_state_dict["layer.lora_A"], mocked_state_dict["layer.lora_A"]) assert torch.equal(filtered_state_dict["layer.lora_B"], mocked_state_dict["layer.lora_B"]) - assert "layer.frozen_weight" not in filtered_state_dict \ No newline at end of file + assert "layer.frozen_weight" not in filtered_state_dict diff --git a/fastvideo/tests/training/distill/test_anyflow_onpolicy.py b/fastvideo/tests/training/distill/test_anyflow_onpolicy.py index 7040a2548..ac4e1a7b4 100644 --- a/fastvideo/tests/training/distill/test_anyflow_onpolicy.py +++ b/fastvideo/tests/training/distill/test_anyflow_onpolicy.py @@ -18,7 +18,6 @@ import torch from fastvideo.train.methods.distribution_matching.anyflow import AnyFlowMethod - # --------------------------------------------------------------------------- # Helpers: build a "naked" AnyFlowMethod that skips __init__. # --------------------------------------------------------------------------- @@ -61,19 +60,16 @@ def _naked_method( def test_get_rollout_schedule_uses_t_list_override_verbatim() -> None: - method = _naked_method( - t_list_override=[999.0, 937.0, 833.0, 624.0, 0.0]) + method = _naked_method(t_list_override=[999.0, 937.0, 833.0, 624.0, 0.0]) schedule = method._get_rollout_schedule(device=torch.device("cpu")) torch.testing.assert_close( schedule, - torch.tensor([999.0, 937.0, 833.0, 624.0, 0.0], - dtype=torch.float32), + torch.tensor([999.0, 937.0, 833.0, 624.0, 0.0], dtype=torch.float32), ) def test_get_rollout_schedule_falls_back_to_denoising_step_list() -> None: - method = _naked_method( - denoising_step_list=[999, 750, 500, 250]) + method = _naked_method(denoising_step_list=[999, 750, 500, 250]) schedule = method._get_rollout_schedule(device=torch.device("cpu")) # Tail must be a 0 boundary so the final Euler step lands at t=0. assert float(schedule[-1].item()) == 0.0 @@ -85,8 +81,7 @@ def test_get_rollout_schedule_falls_back_to_denoising_step_list() -> None: def test_get_rollout_schedule_does_not_double_append_zero() -> None: - method = _naked_method( - denoising_step_list=[999, 750, 500, 0]) + method = _naked_method(denoising_step_list=[999, 750, 500, 0]) schedule = method._get_rollout_schedule(device=torch.device("cpu")) assert schedule.numel() == 4 # no extra boundary inserted. @@ -138,8 +133,7 @@ class _SpyStudent: attn_kind: str = "vsa", ) -> torch.Tensor: del batch, conditional, cfg_uncond, attn_kind - self.seen.append((float(t.flatten()[0].item()), - float(r.flatten()[0].item()))) + self.seen.append((float(t.flatten()[0].item()), float(r.flatten()[0].item()))) # Constant velocity field of magnitude param so we can backprop. return self.param * torch.ones_like(noisy) @@ -195,9 +189,8 @@ def test_rollout_with_grad_true_produces_differentiable_output() -> None: method.student = student # type: ignore[assignment] batch = _make_batch((1, 2, 4, 4, 4)) out = method._student_rollout(batch, with_grad=True) - assert out.requires_grad, ( - "Rollout output must keep a gradient so the DMD loss can backprop " - "through the chosen step.") + assert out.requires_grad, ("Rollout output must keep a gradient so the DMD loss can backprop " + "through the chosen step.") out.sum().backward() assert student.param.grad is not None assert student.param.grad.abs().sum() > 0 @@ -219,13 +212,11 @@ def test_rollout_with_grad_false_blocks_gradient_completely() -> None: def test_broadcast_grad_step_index_in_range() -> None: method = _naked_method(student_sample_steps=4) for _ in range(20): - idx = method._broadcast_grad_step_index( - num_steps=4, device=torch.device("cpu")) + idx = method._broadcast_grad_step_index(num_steps=4, device=torch.device("cpu")) assert 0 <= idx < 4 def test_broadcast_grad_step_index_rejects_non_positive_num_steps() -> None: method = _naked_method() with pytest.raises(ValueError, match="num_steps must be positive"): - method._broadcast_grad_step_index( - num_steps=0, device=torch.device("cpu")) + method._broadcast_grad_step_index(num_steps=0, device=torch.device("cpu")) diff --git a/fastvideo/tests/training/distill/test_anyflow_pretrain.py b/fastvideo/tests/training/distill/test_anyflow_pretrain.py index fee6f5878..8bcfe55c8 100644 --- a/fastvideo/tests/training/distill/test_anyflow_pretrain.py +++ b/fastvideo/tests/training/distill/test_anyflow_pretrain.py @@ -22,7 +22,6 @@ import torch from fastvideo.configs.models.dits import WanVideoConfig - # --------------------------------------------------------------------------- # Task 1: r_embedder config flags default to bit-identity preservation. # --------------------------------------------------------------------------- @@ -96,7 +95,7 @@ def test_embedder_default_path_is_bit_identical_to_legacy() -> None: (no r_timestep kwarg, no extra computation).""" torch.manual_seed(0) emb = _make_embedder(r_embedder=False) - t = torch.randint(0, 1000, (2,), dtype=torch.long) + t = torch.randint(0, 1000, (2, ), dtype=torch.long) txt = torch.randn(2, 4, 16) temb_a, proj_a, _, _ = emb(t, txt) # Calling without r_timestep again must be deterministic-equal. @@ -115,7 +114,7 @@ def test_embedder_enabled_without_r_timestep_is_bit_identical_to_legacy() -> Non emb_legacy = _make_embedder(r_embedder=False) torch.manual_seed(0) emb_dual = _make_embedder(r_embedder=True, fusion="additive") - t = torch.randint(0, 1000, (2,), dtype=torch.long) + t = torch.randint(0, 1000, (2, ), dtype=torch.long) txt = torch.randn(2, 4, 16) temb_legacy, proj_legacy, _, _ = emb_legacy(t, txt) temb_dual, proj_dual, _, _ = emb_dual(t, txt) # No r_timestep. @@ -158,8 +157,7 @@ def test_embedder_additive_fusion_formula() -> None: def test_embedder_deltatime_type_t_minus_r() -> None: """When deltatime_type='t-r', delta_embedder consumes (t - r).""" torch.manual_seed(0) - emb = _make_embedder( - r_embedder=True, fusion="gated", gate=0.5, deltatime_type="t-r") + emb = _make_embedder(r_embedder=True, fusion="gated", gate=0.5, deltatime_type="t-r") t = torch.tensor([800, 800], dtype=torch.long) r = torch.tensor([300, 300], dtype=torch.long) txt = torch.randn(2, 4, 16) @@ -236,9 +234,8 @@ def test_wan_transformer_forward_threads_r_timestep_to_embedder() -> None: from fastvideo.models.dits.wanvideo import WanTransformer3DModel src = inspect.getsource(WanTransformer3DModel.forward) - assert "r_timestep=r_timestep" in src, ( - "WanTransformer3DModel.forward must forward r_timestep into " - "self.condition_embedder") + assert "r_timestep=r_timestep" in src, ("WanTransformer3DModel.forward must forward r_timestep into " + "self.condition_embedder") # --------------------------------------------------------------------------- @@ -249,8 +246,7 @@ def test_wan_transformer_forward_threads_r_timestep_to_embedder() -> None: def _scheduler(*, shift: float = 1.0, n_train: int = 1000): from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import ( FlowMapEulerDiscreteScheduler, ) - return FlowMapEulerDiscreteScheduler( - num_train_timesteps=n_train, shift=shift) + return FlowMapEulerDiscreteScheduler(num_train_timesteps=n_train, shift=shift) def test_flow_map_scheduler_set_timesteps_descending() -> None: @@ -272,8 +268,7 @@ def test_flow_map_scheduler_set_timesteps_custom_overrides_schedule() -> None: device=torch.device("cpu"), custom_timesteps=custom, ) - torch.testing.assert_close( - sched.timesteps, torch.tensor(custom, dtype=torch.float32)) + torch.testing.assert_close(sched.timesteps, torch.tensor(custom, dtype=torch.float32)) def test_flow_map_scheduler_custom_timesteps_must_be_descending() -> None: @@ -292,8 +287,7 @@ def test_flow_map_scheduler_apply_shift_endpoints_invariant() -> None: sched = _scheduler(shift=5.0) t = torch.tensor([0.0, 0.5, 1.0]) shifted = sched.apply_shift(t) - torch.testing.assert_close( - shifted, torch.tensor([0.0, 5.0 / 6.0, 1.0]), rtol=1e-6, atol=1e-6) + torch.testing.assert_close(shifted, torch.tensor([0.0, 5.0 / 6.0, 1.0]), rtol=1e-6, atol=1e-6) def test_flow_map_scheduler_apply_shift_shift_one_is_identity() -> None: @@ -386,8 +380,8 @@ def test_sample_pair_timesteps_partitions_batch_correctly() -> None: device=torch.device("cpu"), generator=None, ) - assert t.shape == (8,) - assert r.shape == (8,) + assert t.shape == (8, ) + assert r.shape == (8, ) assert int(is_diffusion.sum()) == 4 assert int(is_consistency.sum()) == 2 # The masks must be disjoint. @@ -531,12 +525,26 @@ def test_central_difference_dF_dt_guidance_scaling() -> None: t = torch.tensor([500.0]) r = torch.tensor([100.0]) - dF_g1 = _central_difference_dF_dt( - student=student, batch=None, noisy=x, latents=latents, noise=noise, - t=t, r=r, delta=5.0, num_train_timesteps=1000.0, guidance_scale=1.0) - dF_g3 = _central_difference_dF_dt( - student=student, batch=None, noisy=x, latents=latents, noise=noise, - t=t, r=r, delta=5.0, num_train_timesteps=1000.0, guidance_scale=3.0) + dF_g1 = _central_difference_dF_dt(student=student, + batch=None, + noisy=x, + latents=latents, + noise=noise, + t=t, + r=r, + delta=5.0, + num_train_timesteps=1000.0, + guidance_scale=1.0) + dF_g3 = _central_difference_dF_dt(student=student, + batch=None, + noisy=x, + latents=latents, + noise=noise, + t=t, + r=r, + delta=5.0, + num_train_timesteps=1000.0, + guidance_scale=3.0) torch.testing.assert_close(dF_g1 / 3.0, dF_g3, rtol=1e-5, atol=1e-5) @@ -572,9 +580,9 @@ def test_param_names_mapping_includes_delta_embedder_rename() -> None: delta_embedder weights into FastVideo's internal mlp.fc_in/fc_out layout.""" arch = WanVideoConfig().arch_config mapping_keys = list(arch.param_names_mapping.keys()) - assert any("delta_embedder" in k for k in mapping_keys), ( - "WanVideoArchConfig.param_names_mapping must include delta_embedder " - "rename so HF AnyFlow checkpoints load without a separate adapter") + assert any("delta_embedder" in k + for k in mapping_keys), ("WanVideoArchConfig.param_names_mapping must include delta_embedder " + "rename so HF AnyFlow checkpoints load without a separate adapter") def test_param_names_mapping_default_doesnt_break_plain_wan_keys() -> None: @@ -592,13 +600,10 @@ def test_param_names_mapping_default_doesnt_break_plain_wan_keys() -> None: "blocks.0.ffn.net.0.proj.weight", ] arch = WanVideoConfig().arch_config - delta_regexes = [ - k for k in arch.param_names_mapping if "delta_embedder" in k - ] + delta_regexes = [k for k in arch.param_names_mapping if "delta_embedder" in k] assert delta_regexes, "expected at least one delta_embedder regex" for plain in plain_wan_keys: for rx in delta_regexes: - assert re.match(rx, plain) is None, ( - f"delta_embedder regex {rx!r} unexpectedly matched plain " - f"Wan key {plain!r}") + assert re.match(rx, plain) is None, (f"delta_embedder regex {rx!r} unexpectedly matched plain " + f"Wan key {plain!r}") diff --git a/fastvideo/tests/training/distill/test_anyflow_smoke.py b/fastvideo/tests/training/distill/test_anyflow_smoke.py index 876990142..292b1fa8c 100644 --- a/fastvideo/tests/training/distill/test_anyflow_smoke.py +++ b/fastvideo/tests/training/distill/test_anyflow_smoke.py @@ -19,14 +19,9 @@ from pathlib import Path import pytest - REPO_ROOT = Path(__file__).resolve().parents[4] -PRETRAIN_YAML = ( - REPO_ROOT - / "examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml") -ONPOLICY_YAML = ( - REPO_ROOT - / "examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml") +PRETRAIN_YAML = (REPO_ROOT / "examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml") +ONPOLICY_YAML = (REPO_ROOT / "examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml") NUM_NODES = "1" NUM_GPUS_PER_NODE = "2" @@ -45,9 +40,7 @@ def _have_enough_gpus() -> bool: return torch.cuda.device_count() >= 2 -pytestmark = pytest.mark.skipif( - not _have_enough_gpus(), - reason="AnyFlow smoke test requires >= 2 CUDA devices") +pytestmark = pytest.mark.skipif(not _have_enough_gpus(), reason="AnyFlow smoke test requires >= 2 CUDA devices") def _run_torchrun(config_path: Path, *, output_dir: Path) -> None: @@ -64,16 +57,26 @@ def _run_torchrun(config_path: Path, *, output_dir: Path) -> None: sys.executable, "-m", "torch.distributed.run", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE, - "--master_port", env["MASTER_PORT"], - "-m", "fastvideo.train.entrypoint.train", - "--config", str(config_path), - "--training.loop.max_train_steps", "2", - "--training.checkpoint.output_dir", str(output_dir), - "--training.distributed.num_gpus", NUM_GPUS_PER_NODE, - "--training.distributed.hsdp_shard_dim", NUM_GPUS_PER_NODE, - "--training.data.train_batch_size", "1", + "--nnodes", + NUM_NODES, + "--nproc_per_node", + NUM_GPUS_PER_NODE, + "--master_port", + env["MASTER_PORT"], + "-m", + "fastvideo.train.entrypoint.train", + "--config", + str(config_path), + "--training.loop.max_train_steps", + "2", + "--training.checkpoint.output_dir", + str(output_dir), + "--training.distributed.num_gpus", + NUM_GPUS_PER_NODE, + "--training.distributed.hsdp_shard_dim", + NUM_GPUS_PER_NODE, + "--training.data.train_batch_size", + "1", ] process = subprocess.run(cmd, capture_output=True, text=True, env=env) if process.stdout: @@ -81,8 +84,7 @@ def _run_torchrun(config_path: Path, *, output_dir: Path) -> None: if process.stderr: print("STDERR:", process.stderr) if process.returncode != 0: - raise subprocess.CalledProcessError( - process.returncode, cmd, process.stdout, process.stderr) + raise subprocess.CalledProcessError(process.returncode, cmd, process.stdout, process.stderr) def test_anyflow_pretrain_smoke(tmp_path: Path) -> None: @@ -107,18 +109,30 @@ def test_anyflow_onpolicy_smoke(tmp_path: Path) -> None: sys.executable, "-m", "torch.distributed.run", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE, - "--master_port", env["MASTER_PORT"], - "-m", "fastvideo.train.entrypoint.train", - "--config", str(ONPOLICY_YAML), - "--training.loop.max_train_steps", "2", - "--training.checkpoint.output_dir", str(tmp_path / "onpolicy"), - "--training.distributed.num_gpus", NUM_GPUS_PER_NODE, - "--training.distributed.hsdp_shard_dim", NUM_GPUS_PER_NODE, - "--training.data.train_batch_size", "1", - "--models.student.init_from", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--method.student_sample_steps", "2", + "--nnodes", + NUM_NODES, + "--nproc_per_node", + NUM_GPUS_PER_NODE, + "--master_port", + env["MASTER_PORT"], + "-m", + "fastvideo.train.entrypoint.train", + "--config", + str(ONPOLICY_YAML), + "--training.loop.max_train_steps", + "2", + "--training.checkpoint.output_dir", + str(tmp_path / "onpolicy"), + "--training.distributed.num_gpus", + NUM_GPUS_PER_NODE, + "--training.distributed.hsdp_shard_dim", + NUM_GPUS_PER_NODE, + "--training.data.train_batch_size", + "1", + "--models.student.init_from", + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "--method.student_sample_steps", + "2", ] process = subprocess.run(cmd, capture_output=True, text=True, env=env) if process.stdout: @@ -126,5 +140,4 @@ def test_anyflow_onpolicy_smoke(tmp_path: Path) -> None: if process.stderr: print("STDERR:", process.stderr) if process.returncode != 0: - raise subprocess.CalledProcessError( - process.returncode, cmd, process.stdout, process.stderr) + raise subprocess.CalledProcessError(process.returncode, cmd, process.stdout, process.stderr) diff --git a/fastvideo/tests/training/distill/test_distill_dmd.py b/fastvideo/tests/training/distill/test_distill_dmd.py index 84504e2c5..806780180 100644 --- a/fastvideo/tests/training/distill/test_distill_dmd.py +++ b/fastvideo/tests/training/distill/test_distill_dmd.py @@ -1,4 +1,5 @@ import os + os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29512" import sys @@ -27,100 +28,69 @@ def run_worker(): parser = FlexibleArgumentParser() parser = TrainingArgs.add_cli_args(parser) parser = FastVideoArgs.add_cli_args(parser) - + # Set the arguments as they are in finetune_v1_test.sh args = parser.parse_args([ - "--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--inference_mode", "False", - "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--real_score_model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--fake_score_model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--data_path", "data/crush-smol_processed_t2v/combined_parquet_dataset", - "--validation_dataset_file", "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json", - "--train_batch_size", "1", - "--num_latent_t", "4", - "--num_gpus", "2", - "--sp_size", "2", - "--tp_size", "1", - "--hsdp_replicate_dim", "1", - "--hsdp_shard_dim", "2", - "--train_sp_batch_size", "1", - "--dataloader_num_workers", "1", - "--gradient_accumulation_steps", "2", - "--max_train_steps", "2", - "--learning_rate", "1e-5", - "--mixed_precision", "bf16", - "--training_state_checkpointing_steps", "30", - "--weight_only_checkpointing_steps", "30", - "--validation_steps", "10", - "--validation_sampling_steps", "3", - "--log_validation", - "--checkpoints_total_limit", "3", - "--ema_start_step", "0", - "--training_cfg_rate", "0.0", - "--output_dir", "data/wan_finetune_test", - "--tracker_project_name", "wan_finetune_ci", - "--wandb_run_name", wandb_name, - "--num_height", "480", - "--num_width", "832", - "--num_frames", "13", - "--flow_shift", "8", - "--weight_decay", "0.01", - "--dit_precision", "fp32", - "--max_grad_norm", "1.0", - "--generator_update_interval", "5", - "--dmd_denoising_steps", "1000,757,522", - "--min_timestep_ratio", "0.02", - "--max_timestep_ratio", "0.98", - "--real_score_guidance_scale", "3.5", + "--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--inference_mode", "False", + "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--real_score_model_path", + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--fake_score_model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "--data_path", "data/crush-smol_processed_t2v/combined_parquet_dataset", "--validation_dataset_file", + "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json", "--train_batch_size", "1", + "--num_latent_t", "4", "--num_gpus", "2", "--sp_size", "2", "--tp_size", "1", "--hsdp_replicate_dim", "1", + "--hsdp_shard_dim", "2", "--train_sp_batch_size", "1", "--dataloader_num_workers", "1", + "--gradient_accumulation_steps", "2", "--max_train_steps", "2", "--learning_rate", "1e-5", "--mixed_precision", + "bf16", "--training_state_checkpointing_steps", "30", "--weight_only_checkpointing_steps", "30", + "--validation_steps", "10", "--validation_sampling_steps", "3", "--log_validation", "--checkpoints_total_limit", + "3", "--ema_start_step", "0", "--training_cfg_rate", "0.0", "--output_dir", "data/wan_finetune_test", + "--tracker_project_name", "wan_finetune_ci", "--wandb_run_name", wandb_name, "--num_height", "480", + "--num_width", "832", "--num_frames", "13", "--flow_shift", "8", "--weight_decay", "0.01", "--dit_precision", + "fp32", "--max_grad_norm", "1.0", "--generator_update_interval", "5", "--dmd_denoising_steps", "1000,757,522", + "--min_timestep_ratio", "0.02", "--max_timestep_ratio", "0.98", "--real_score_guidance_scale", "3.5", "--enable_gradient_checkpointing_type", "full" ]) # Call the main training function - pipeline = WanDistillationPipeline.from_pretrained( - args.pretrained_model_name_or_path, args=args) + pipeline = WanDistillationPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args) args = pipeline.training_args pipeline.train() logger.info("Training pipeline done") + def test_distributed_training(): """Test the distributed training setup""" os.environ["WANDB_MODE"] = "offline" data_dir = Path("data/crush-smol_processed_t2v") - + if not data_dir.exists(): print(f"Downloading test dataset to {data_dir}...") - snapshot_download( - repo_id="wlsaidhi/crush-smol_processed_t2v", - local_dir=str(data_dir), - repo_type="dataset", - local_dir_use_symlinks=False - ) - + snapshot_download(repo_id="wlsaidhi/crush-smol_processed_t2v", + local_dir=str(data_dir), + repo_type="dataset", + local_dir_use_symlinks=False) + # Get the current file path current_file = Path(__file__).resolve() - + # Run torchrun command cmd = [ - "torchrun", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE, - "--master_port", os.environ["MASTER_PORT"], + "torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, "--master_port", + os.environ["MASTER_PORT"], str(current_file) ] process = subprocess.run(cmd, capture_output=True, text=True) - + # Print stdout and stderr for debugging if process.stdout: print("STDOUT:", process.stdout) if process.stderr: print("STDERR:", process.stderr) - + # Check if the process failed if process.returncode != 0: print(f"Process failed with return code: {process.returncode}") raise subprocess.CalledProcessError(process.returncode, cmd, process.stdout, process.stderr) + if __name__ == "__main__": if os.environ.get("LOCAL_RANK") is not None: # We're being run by torchrun diff --git a/fastvideo/tests/training/lora/test_lora_training.py b/fastvideo/tests/training/lora/test_lora_training.py index 9649fdc02..bf8edceb0 100644 --- a/fastvideo/tests/training/lora/test_lora_training.py +++ b/fastvideo/tests/training/lora/test_lora_training.py @@ -1,4 +1,5 @@ import os + os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29514" import sys @@ -20,67 +21,35 @@ def test_lora_training(): os.environ["WANDB_MODE"] = "online" data_dir = Path("data/crush-smol_processed_t2v") - + if not data_dir.exists(): print(f"Downloading test dataset to {data_dir}...") - snapshot_download( - repo_id="wlsaidhi/crush-smol_processed_t2v", - local_dir=str(data_dir), - repo_type="dataset", - local_dir_use_symlinks=False - ) + snapshot_download(repo_id="wlsaidhi/crush-smol_processed_t2v", + local_dir=str(data_dir), + repo_type="dataset", + local_dir_use_symlinks=False) # Run torchrun command directly on the training pipeline like the shell script cmd = [ - "torchrun", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE, - "--master_port", os.environ["MASTER_PORT"], - "fastvideo/training/wan_training_pipeline.py", - "--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--inference_mode", "False", - "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--data_path", "data/crush-smol_processed_t2v/combined_parquet_dataset", + "torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, "--master_port", + os.environ["MASTER_PORT"], "fastvideo/training/wan_training_pipeline.py", "--model_path", + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--inference_mode", "False", "--pretrained_model_name_or_path", + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "--data_path", "data/crush-smol_processed_t2v/combined_parquet_dataset", "--validation_dataset_file", "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json", - "--train_batch_size", "1", - "--num_latent_t", "8", - "--num_gpus", NUM_GPUS_PER_NODE, - "--sp_size", NUM_GPUS_PER_NODE, - "--tp_size", NUM_GPUS_PER_NODE, - "--hsdp_replicate_dim", "1", - "--hsdp_shard_dim", "2", - "--train_sp_batch_size", "1", - "--dataloader_num_workers", "1", - "--gradient_accumulation_steps", "8", - "--max_train_steps", "5", - "--learning_rate", "5e-5", - "--mixed_precision", "bf16", - "--weight_only_checkpointing_steps", "6000", - "--training_state_checkpointing_steps", "6000", - "--validation_steps", "50", - "--validation_sampling_steps", "50", - "--log_validation", - "--checkpoints_total_limit", "3", - "--ema_start_step", "0", - "--training_cfg_rate", "0.1", - "--output_dir", "/workspace", - "--tracker_project_name", "wan_lora_finetune_ci", - "--wandb_run_name", wandb_name, - "--num_height", "480", - "--num_width", "832", - "--num_frames", "77", - "--validation_guidance_scale", "1.0", - "--num_euler_timesteps", "50", - "--multi_phased_distill_schedule", "4000-1", - "--weight_decay", "1e-4", - "--not_apply_cfg_solver", - "--dit_precision", "fp32", - "--max_grad_norm", "1.0", - "--lora_rank", "32", - "--lora_training", "True", - "--seed", "42" + "--train_batch_size", "1", "--num_latent_t", "8", "--num_gpus", NUM_GPUS_PER_NODE, "--sp_size", + NUM_GPUS_PER_NODE, "--tp_size", NUM_GPUS_PER_NODE, "--hsdp_replicate_dim", "1", "--hsdp_shard_dim", "2", + "--train_sp_batch_size", "1", "--dataloader_num_workers", "1", "--gradient_accumulation_steps", "8", + "--max_train_steps", "5", "--learning_rate", "5e-5", "--mixed_precision", "bf16", + "--weight_only_checkpointing_steps", "6000", "--training_state_checkpointing_steps", "6000", + "--validation_steps", "50", "--validation_sampling_steps", "50", "--log_validation", + "--checkpoints_total_limit", "3", "--ema_start_step", "0", "--training_cfg_rate", "0.1", "--output_dir", + "/workspace", "--tracker_project_name", "wan_lora_finetune_ci", "--wandb_run_name", wandb_name, "--num_height", + "480", "--num_width", "832", "--num_frames", "77", "--validation_guidance_scale", "1.0", + "--num_euler_timesteps", "50", "--multi_phased_distill_schedule", "4000-1", "--weight_decay", "1e-4", + "--not_apply_cfg_solver", "--dit_precision", "fp32", "--max_grad_norm", "1.0", "--lora_rank", "32", + "--lora_training", "True", "--seed", "42" ] - + process = subprocess.run(cmd, check=True) summary_file = '/workspace/tracker/wandb/latest-run/files/wandb-summary.json' @@ -96,8 +65,8 @@ def test_lora_training(): fields_and_thresholds = { 'avg_step_time': 30.0, # something up with modal # 'grad_norm': 0.05, # too volatile for now. TODO: fix nondeterminism in training - 'step_time': 30.0, # something up with modal - 'train_loss': 0.05 + 'step_time': 30.0, # something up with modal + 'train_loss': 0.05 } failures = [] @@ -106,9 +75,13 @@ def test_lora_training(): ref_value = reference_wandb_summary[field] current_value = wandb_summary[field] diff = abs(ref_value - current_value) - print(f"INFO: {field}, diff: {diff}, threshold: {threshold}, reference: {ref_value}, current: {current_value}") + print( + f"INFO: {field}, diff: {diff}, threshold: {threshold}, reference: {ref_value}, current: {current_value}" + ) if diff > threshold: - failures.append(f"FAILED: {field} difference {diff} exceeds threshold of {threshold} (reference: {ref_value}, current: {current_value})") + failures.append( + f"FAILED: {field} difference {diff} exceeds threshold of {threshold} (reference: {ref_value}, current: {current_value})" + ) else: print(f"WARNING: Field {field} not found in one or both summary files") @@ -117,4 +90,4 @@ def test_lora_training(): if __name__ == "__main__": - test_lora_training() \ No newline at end of file + test_lora_training() diff --git a/fastvideo/tests/training/self-forcing/test_self_forcing.py b/fastvideo/tests/training/self-forcing/test_self_forcing.py index a9f84fa82..db42c2473 100644 --- a/fastvideo/tests/training/self-forcing/test_self_forcing.py +++ b/fastvideo/tests/training/self-forcing/test_self_forcing.py @@ -1,4 +1,5 @@ import os + os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29513" import sys @@ -26,118 +27,169 @@ def run_worker(): parser = FlexibleArgumentParser() parser = TrainingArgs.add_cli_args(parser) parser = FastVideoArgs.add_cli_args(parser) - + # Set the arguments based on the distill_dmd_t2v_1.3B.sh script args = parser.parse_args([ - "--model_path", "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", - "--inference_mode", "False", - "--pretrained_model_name_or_path", "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", - "--real_score_model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--fake_score_model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "--data_path", "data/crush-smol_processed_t2v/combined_parquet_dataset", - "--validation_dataset_file", "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json", - "--train_batch_size", "1", - "--num_latent_t", "21", - "--num_gpus", "2", - "--sp_size", "1", - "--tp_size", "1", - "--hsdp_replicate_dim", "1", - "--hsdp_shard_dim", "2", - "--train_sp_batch_size", "1", - "--dataloader_num_workers", "1", - "--gradient_accumulation_steps", "1", - "--max_train_steps", "2", - "--learning_rate", "1e-5", - "--mixed_precision", "bf16", - "--training_state_checkpointing_steps", "30", - "--weight_only_checkpointing_steps", "30", - "--validation_steps", "10", - "--validation_sampling_steps", "3", + "--model_path", + "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", + "--inference_mode", + "False", + "--pretrained_model_name_or_path", + "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", + "--real_score_model_path", + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "--fake_score_model_path", + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "--data_path", + "data/crush-smol_processed_t2v/combined_parquet_dataset", + "--validation_dataset_file", + "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json", + "--train_batch_size", + "1", + "--num_latent_t", + "21", + "--num_gpus", + "2", + "--sp_size", + "1", + "--tp_size", + "1", + "--hsdp_replicate_dim", + "1", + "--hsdp_shard_dim", + "2", + "--train_sp_batch_size", + "1", + "--dataloader_num_workers", + "1", + "--gradient_accumulation_steps", + "1", + "--max_train_steps", + "2", + "--learning_rate", + "1e-5", + "--mixed_precision", + "bf16", + "--training_state_checkpointing_steps", + "30", + "--weight_only_checkpointing_steps", + "30", + "--validation_steps", + "10", + "--validation_sampling_steps", + "3", "--log_validation", - "--checkpoints_total_limit", "3", - "--ema_start_step", "0", - "--training_cfg_rate", "0.0", - "--output_dir", "data/wan_self_forcing_test", - "--tracker_project_name", "wan_self_forcing_ci", - "--wandb_run_name", wandb_name, - "--num_height", "480", - "--num_width", "832", - "--num_frames", "21", - "--flow_shift", "5", - "--weight_decay", "0.01", - "--dit_precision", "fp32", - "--max_grad_norm", "1.0", + "--checkpoints_total_limit", + "3", + "--ema_start_step", + "0", + "--training_cfg_rate", + "0.0", + "--output_dir", + "data/wan_self_forcing_test", + "--tracker_project_name", + "wan_self_forcing_ci", + "--wandb_run_name", + wandb_name, + "--num_height", + "480", + "--num_width", + "832", + "--num_frames", + "21", + "--flow_shift", + "5", + "--weight_decay", + "0.01", + "--dit_precision", + "fp32", + "--max_grad_norm", + "1.0", # DMD args - "--dmd_denoising_steps", "1000,750,500", # Reduced steps for testing - "--min_timestep_ratio", "0.02", - "--max_timestep_ratio", "0.98", - "--dfake_gen_update_ratio", "5", - "--real_score_guidance_scale", "3.0", - "--fake_score_learning_rate", "8e-6", - "--fake_score_betas", "0.0,0.999", + "--dmd_denoising_steps", + "1000,750,500", # Reduced steps for testing + "--min_timestep_ratio", + "0.02", + "--max_timestep_ratio", + "0.98", + "--dfake_gen_update_ratio", + "5", + "--real_score_guidance_scale", + "3.0", + "--fake_score_learning_rate", + "8e-6", + "--fake_score_betas", + "0.0,0.999", "--warp_denoising_step", - "--enable_gradient_checkpointing_type", "full", + "--enable_gradient_checkpointing_type", + "full", # Self-forcing specific args "--log_visualization", "--simulate_generator_forward", - "--num_frame_per_block", "3", + "--num_frame_per_block", + "3", "--enable_gradient_masking", - "--gradient_mask_last_n_frames", "21", - "--independent_first_frame", "False", - "--same_step_across_blocks", "True", - "--last_step_only", "False", - "--context_noise", "0", - "--use_ema", "True", - "--ema_decay", "0.99", - "--ema_start_step", "100", + "--gradient_mask_last_n_frames", + "21", + "--independent_first_frame", + "False", + "--same_step_across_blocks", + "True", + "--last_step_only", + "False", + "--context_noise", + "0", + "--use_ema", + "True", + "--ema_decay", + "0.99", + "--ema_start_step", + "100", ]) - + # Call the main training function - pipeline = WanSelfForcingDistillationPipeline.from_pretrained( - args.pretrained_model_name_or_path, args=args) + pipeline = WanSelfForcingDistillationPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args) args = pipeline.training_args pipeline.train() logger.info("Self-forcing distillation training pipeline done") + def test_distributed_training(): """Test the distributed self-forcing training setup""" os.environ["WANDB_MODE"] = "offline" data_dir = Path("data/crush-smol_processed_t2v") - + if not data_dir.exists(): print(f"Downloading test dataset to {data_dir}...") - snapshot_download( - repo_id="wlsaidhi/crush-smol_processed_t2v", - local_dir=str(data_dir), - repo_type="dataset", - local_dir_use_symlinks=False - ) - + snapshot_download(repo_id="wlsaidhi/crush-smol_processed_t2v", + local_dir=str(data_dir), + repo_type="dataset", + local_dir_use_symlinks=False) + # Get the current file path current_file = Path(__file__).resolve() - + # Run torchrun command cmd = [ - "torchrun", - "--nnodes", NUM_NODES, - "--nproc_per_node", NUM_GPUS_PER_NODE, - "--master_port", os.environ["MASTER_PORT"], + "torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, "--master_port", + os.environ["MASTER_PORT"], str(current_file) ] process = subprocess.run(cmd, capture_output=True, text=True) - + # Print stdout and stderr for debugging if process.stdout: print("STDOUT:", process.stdout) if process.stderr: print("STDERR:", process.stderr) - + # Check if the process failed if process.returncode != 0: print(f"Process failed with return code: {process.returncode}") raise subprocess.CalledProcessError(process.returncode, cmd, process.stdout, process.stderr) + if __name__ == "__main__": if os.environ.get("LOCAL_RANK") is not None: # We're being run by torchrun @@ -145,4 +197,3 @@ if __name__ == "__main__": else: # We're being run directly test_distributed_training() - diff --git a/fastvideo/tests/training/test_trackers.py b/fastvideo/tests/training/test_trackers.py index 95bf01f17..cbfc503aa 100644 --- a/fastvideo/tests/training/test_trackers.py +++ b/fastvideo/tests/training/test_trackers.py @@ -23,6 +23,7 @@ from fastvideo.training.trackers import ( class _FakeSwanlab(ModuleType): + def __init__(self) -> None: super().__init__("swanlab") self.init_kwargs: dict[str, Any] | None = None @@ -194,6 +195,7 @@ def test_swanlab_missing_dependency_error_is_actionable( class _RecordingTracker(BaseTracker): + def __init__(self, name: str) -> None: super().__init__() self.name = name diff --git a/fastvideo/tests/transformers/test_cosmos.py b/fastvideo/tests/transformers/test_cosmos.py index e5629ff38..50930dec4 100644 --- a/fastvideo/tests/transformers/test_cosmos.py +++ b/fastvideo/tests/transformers/test_cosmos.py @@ -16,7 +16,6 @@ from fastvideo.utils import maybe_download_model from fastvideo.configs.models.dits import CosmosVideoConfig from fastvideo.pipelines.pipeline_batch_info import ForwardBatch - logger = init_logger(__name__) os.environ["MASTER_ADDR"] = "localhost" @@ -51,13 +50,13 @@ def test_cosmos2_transformer(): loader = TransformerLoader() model2 = loader.load(TRANSFORMER_PATH, args).to(device, dtype=precision) - model1 = CosmosTransformer3DModel.from_pretrained( - TRANSFORMER_PATH, torch_dtype=precision).to(device, dtype=precision).requires_grad_(False) + model1 = CosmosTransformer3DModel.from_pretrained(TRANSFORMER_PATH, + torch_dtype=precision).to(device, + dtype=precision).requires_grad_(False) total_params = sum(p.numel() for p in model1.parameters()) # Calculate weight sum for model1 (converting to float64 to avoid overflow) - weight_sum_model1 = sum( - p.to(torch.float64).sum().item() for p in model1.parameters()) + weight_sum_model1 = sum(p.to(torch.float64).sum().item() for p in model1.parameters()) # Also calculate mean for more stable comparison weight_mean_model1 = weight_sum_model1 / total_params logger.info("Model 1 weight sum: %s", weight_sum_model1) @@ -65,8 +64,7 @@ def test_cosmos2_transformer(): # Calculate weight sum for model2 (converting to float64 to avoid overflow) total_params_model2 = sum(p.numel() for p in model2.parameters()) - weight_sum_model2 = sum( - p.to(torch.float64).sum().item() for p in model2.parameters()) + weight_sum_model2 = sum(p.to(torch.float64).sum().item() for p in model2.parameters()) # Also calculate mean for more stable comparison weight_mean_model2 = weight_sum_model2 / total_params_model2 logger.info("Model 2 weight sum: %s", weight_sum_model2) @@ -86,20 +84,22 @@ def test_cosmos2_transformer(): seq_len = 30 # Video latents [B, C, T, H, W] - Cosmos2 specific dimensions - hidden_states = torch.randn(batch_size, - 17, - 1, # Single frame for image generation - 32, # Height patches - 32, # Width patches - device=device, - dtype=precision) + hidden_states = torch.randn( + batch_size, + 17, + 1, # Single frame for image generation + 32, # Height patches + 32, # Width patches + device=device, + dtype=precision) # Text embeddings [B, L, D] - Cosmos2 uses T5 embeddings with 1024 dim - encoder_hidden_states = torch.randn(batch_size, - seq_len, - 1024, # T5 embedding dimension - device=device, - dtype=precision) + encoder_hidden_states = torch.randn( + batch_size, + seq_len, + 1024, # T5 embedding dimension + device=device, + dtype=precision) # Timestep timestep = torch.tensor([500], device=device, dtype=precision) @@ -108,9 +108,7 @@ def test_cosmos2_transformer(): padding_mask = hidden_states.new_zeros(1, 1, 32, 32, device=device, dtype=precision) # print(padding_mask.shape) - forward_batch = ForwardBatch( - data_type="dummy", - ) + forward_batch = ForwardBatch(data_type="dummy", ) with torch.autocast('cuda', dtype=precision): output1 = model1( diff --git a/fastvideo/tests/transformers/test_cosmos2_5.py b/fastvideo/tests/transformers/test_cosmos2_5.py index 1883e6269..a8d863915 100644 --- a/fastvideo/tests/transformers/test_cosmos2_5.py +++ b/fastvideo/tests/transformers/test_cosmos2_5.py @@ -25,7 +25,6 @@ from fastvideo.utils import maybe_download_model from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig from fastvideo.pipelines.pipeline_batch_info import ForwardBatch - logger = init_logger(__name__) # Log the cosmos-predict2.5 path after logger is initialized @@ -45,9 +44,7 @@ CHECKPOINT_FILENAME = "81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt" def _resolve_model_path() -> str: - skip_if_gated_repo_inaccessible(BASE_MODEL_PATH, - test_name="Cosmos 2.5 transformer test", - allow_module_level=True) + skip_if_gated_repo_inaccessible(BASE_MODEL_PATH, test_name="Cosmos 2.5 transformer test", allow_module_level=True) return maybe_download_model(BASE_MODEL_PATH, local_dir=None) @@ -66,13 +63,13 @@ def load_reference_cosmos25_model(checkpoint_path: str, device, dtype): try: # Try to import from cosmos-predict2.5 repo from cosmos_predict2._src.predict2.networks.minimal_v1_lvg_dit import MinimalV1LVGDiT - + # COSMOS 2.5 2B model configuration model_config = { 'max_img_h': 240, 'max_img_w': 240, 'max_frames': 128, - 'in_channels': 16, + 'in_channels': 16, 'out_channels': 16, 'patch_spatial': 2, 'patch_temporal': 1, @@ -91,19 +88,19 @@ def load_reference_cosmos25_model(checkpoint_path: str, device, dtype): 'rope_t_extrapolation_ratio': 1.0, 'extra_per_block_abs_pos_emb': False, 'rope_enable_fps_modulation': False, - 'use_crossattn_projection': True, - 'crossattn_proj_in_channels': 100352, + 'use_crossattn_projection': True, + 'crossattn_proj_in_channels': 100352, 'concat_padding_mask': True, 'atten_backend': 'torch', } - + model = MinimalV1LVGDiT(**model_config) - + # Load checkpoint if path exists if os.path.exists(checkpoint_path): logger.info(f"Loading reference model from {checkpoint_path}") checkpoint = torch.load(checkpoint_path, map_location='cpu') - + # Extract state dict if 'state_dict' in checkpoint: checkpoint_state = checkpoint['state_dict'] @@ -111,36 +108,35 @@ def load_reference_cosmos25_model(checkpoint_path: str, device, dtype): checkpoint_state = checkpoint['model'] else: checkpoint_state = checkpoint - + # Filter to only model parameters (remove training metadata) - model_state = {k: v for k, v in checkpoint_state.items() - if k.startswith('net.') and 'accum_' not in k} - + model_state = {k: v for k, v in checkpoint_state.items() if k.startswith('net.') and 'accum_' not in k} + # Transform checkpoint keys to match model's expected format # 1. Strip 'net.' prefix (e.g., 'net.blocks.0.self_attn.*' -> 'blocks.0.self_attn.*') # 2. Add '_checkpoint_wrapped_module' after 'blocks.N.' if model expects it transformed_state = {} - + # First, check what the model expects model_state_dict = model.state_dict() needs_checkpoint_wrapper = any('_checkpoint_wrapped_module' in k for k in model_state_dict.keys()) - + for key, value in model_state.items(): # Strip 'net.' prefix new_key = key[4:] if key.startswith('net.') else key # Strip 'net.' prefix - + # Add '_checkpoint_wrapped_module' if needed if needs_checkpoint_wrapper and new_key.startswith('blocks.'): # Pattern: 'blocks.N.something' -> 'blocks.N._checkpoint_wrapped_module.something' parts = new_key.split('.', 2) if len(parts) >= 3 and parts[0] == 'blocks' and parts[1].isdigit(): new_key = f"{parts[0]}.{parts[1]}._checkpoint_wrapped_module.{parts[2]}" - + transformed_state[new_key] = value - + # Load with strict=False to handle any remaining mismatches missing_keys, unexpected_keys = model.load_state_dict(transformed_state, strict=False) - + if missing_keys: logger.warning(f"Missing keys when loading reference model: {len(missing_keys)} keys") # Show all missing keys for debugging @@ -148,30 +144,31 @@ def load_reference_cosmos25_model(checkpoint_path: str, device, dtype): for k in missing_keys: logger.warning(f" - {k}") # Filter out _extra_state and pos_embedder keys as they're optional - missing_important = [k for k in missing_keys - if '_extra_state' not in k and 'pos_embedder' not in k and 'accum_' not in k] + missing_important = [ + k for k in missing_keys if '_extra_state' not in k and 'pos_embedder' not in k and 'accum_' not in k + ] if missing_important: logger.warning(f"Missing important keys ({len(missing_important)} total):") for k in missing_important[:10]: # Show first 10 logger.warning(f" - {k}") if len(missing_important) > 10: logger.warning(f" ... and {len(missing_important) - 10} more") - + if unexpected_keys: logger.warning(f"Unexpected keys when loading reference model: {len(unexpected_keys)} keys") logger.warning("All unexpected keys:") for k in unexpected_keys: logger.warning(f" - {k}") - + logger.info(f"Successfully loaded {len(transformed_state)} parameters into reference model") else: logger.warning(f"Checkpoint path {checkpoint_path} not found, using random weights") - + model = model.to(device, dtype=dtype) model.eval() - + return model - + except ImportError as e: logger.error(f"Failed to import cosmos-predict2.5: {e}") logger.info("Make sure cosmos-predict2.5 is in your Python path") @@ -183,10 +180,10 @@ def test_cosmos25_transformer(): """Test COSMOS 2.5 transformer against reference implementation.""" device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") precision = torch.bfloat16 - + # Create COSMOS 2.5 specific config from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25ArchConfig - + arch_config = Cosmos25ArchConfig( num_attention_heads=16, attention_head_dim=128, # 2048 / 16 @@ -201,18 +198,18 @@ def test_cosmos25_transformer(): adaln_lora_dim=256, use_adaln_lora=True, concat_padding_mask=True, - extra_pos_embed_type=None, - use_crossattn_projection=True, + extra_pos_embed_type=None, + use_crossattn_projection=True, rope_enable_fps_modulation=False, qk_norm="rms_norm", ) - + cosmos25_config = Cosmos25VideoConfig(arch_config=arch_config) - + # Create FastVideo model directly (Cosmos 2.5 is not in diffusers format) logger.info("Creating FastVideo COSMOS 2.5 model...") from fastvideo.models.dits.cosmos2_5 import Cosmos25Transformer3DModel - + # Get hf_config from the arch_config for model initialization hf_config = { 'in_channels': arch_config.in_channels, @@ -233,26 +230,26 @@ def test_cosmos25_transformer(): 'rope_enable_fps_modulation': arch_config.rope_enable_fps_modulation, 'qk_norm': arch_config.qk_norm, } - + fastvideo_model = Cosmos25Transformer3DModel(config=cosmos25_config, hf_config=hf_config) fastvideo_model = fastvideo_model.to(device, dtype=precision) fastvideo_model.eval() - + # Construct checkpoint path using relative paths checkpoint_file = os.path.join(MODEL_PATH, CHECKPOINT_SUBDIR, CHECKPOINT_FILENAME) - + if not os.path.exists(checkpoint_file): logger.warning(f"Checkpoint file not found at {checkpoint_file}") logger.info("Will test architecture without loading checkpoint weights") checkpoint_file = None - + # Load checkpoint into FastVideo model using param_names_mapping if checkpoint_file: logger.info(f"Loading checkpoint into FastVideo model from {checkpoint_file}") from fastvideo.models.loader.utils import hf_to_custom_state_dict, get_param_names_mapping - + checkpoint = torch.load(checkpoint_file, map_location='cpu') - + # Extract state dict (checkpoint might have 'state_dict', 'model', or be the dict itself) if 'state_dict' in checkpoint: checkpoint_state = checkpoint['state_dict'] @@ -260,127 +257,96 @@ def test_cosmos25_transformer(): checkpoint_state = checkpoint['model'] else: checkpoint_state = checkpoint - + # Filter to only model parameters (remove training metadata like accum_*) - model_state = {k: v for k, v in checkpoint_state.items() - if k.startswith('net.') and 'accum_' not in k} - + model_state = {k: v for k, v in checkpoint_state.items() if k.startswith('net.') and 'accum_' not in k} + # Convert checkpoint keys to FastVideo format using param_names_mapping - param_names_mapping_fn = get_param_names_mapping( - cosmos25_config.arch_config.param_names_mapping - ) - custom_state_dict, reverse_mapping = hf_to_custom_state_dict( - model_state, param_names_mapping_fn - ) - + param_names_mapping_fn = get_param_names_mapping(cosmos25_config.arch_config.param_names_mapping) + custom_state_dict, reverse_mapping = hf_to_custom_state_dict(model_state, param_names_mapping_fn) + # Only load keys that exist in the model model_param_names = set(fastvideo_model.state_dict().keys()) filtered_state_dict = { - k: v.to(device=device, dtype=precision) - for k, v in custom_state_dict.items() - if k in model_param_names + k: v.to(device=device, dtype=precision) + for k, v in custom_state_dict.items() if k in model_param_names } - + # Load into FastVideo model - missing_keys, unexpected_keys = fastvideo_model.load_state_dict( - filtered_state_dict, strict=False - ) - + missing_keys, unexpected_keys = fastvideo_model.load_state_dict(filtered_state_dict, strict=False) + if missing_keys: logger.warning(f"Missing keys when loading checkpoint: {len(missing_keys)} keys") # Filter out _extra_state keys as they're optional missing_non_extra = [k for k in missing_keys if '_extra_state' not in k] if missing_non_extra: logger.warning(f"Missing non-extra keys (first 10): {missing_non_extra[:10]}") - + if unexpected_keys: logger.warning(f"Unexpected keys when loading checkpoint: {len(unexpected_keys)} keys") - + logger.info(f"Successfully loaded {len(filtered_state_dict)} parameters into FastVideo model") - + # Try to load reference model from the raw checkpoint logger.info("Loading reference COSMOS 2.5 model...") reference_model = load_reference_cosmos25_model(checkpoint_file, device, precision) if checkpoint_file else None - + # Set models to eval mode fastvideo_model = fastvideo_model.eval() if reference_model is not None: reference_model = reference_model.eval() - + # Create test inputs batch_size = 1 seq_len = 77 # Typical T5 sequence length - + # Video latents [B, C, T, H, W] # COSMOS 2.5: 16 channels (VAE latent), no condition mask in input hidden_states = torch.randn( batch_size, 16, # VAE channels only (condition mask added internally) - 1, # Single frame for image generation (or 16 for video) + 1, # Single frame for image generation (or 16 for video) 64, # Height (720p / 8 / 2 patch = 45, use 64 for testing) 64, # Width device=device, - dtype=precision - ) - + dtype=precision) + # Condition mask [B, 1, T, H, W] - for video2world conditioning - condition_mask = torch.zeros( - batch_size, - 1, - 1, - 64, - 64, - device=device, - dtype=precision - ) - + condition_mask = torch.zeros(batch_size, 1, 1, 64, 64, device=device, dtype=precision) + # Text embeddings [B, L, D] - Qwen 7B embeddings (100,352 dims) # Using 100,352 dimensions to match the crossattn_projection layer - encoder_hidden_states = torch.randn( - batch_size, - seq_len, - 100352, - device=device, - dtype=precision - ) - + encoder_hidden_states = torch.randn(batch_size, seq_len, 100352, device=device, dtype=precision) + # Timestep [B, T] - official model expects [B, T] shape with dtype matching model precision # For single frame, use [B, 1] timestep = torch.full((batch_size, 1), 500.0, device=device, dtype=precision) - + # Padding mask [B, H, W] - official model expects NO channel dimension # It's added internally via unsqueeze(1) if needed - padding_mask = torch.ones( - batch_size, - 64, - 64, - device=device, - dtype=precision - ) - + padding_mask = torch.ones(batch_size, 64, 64, device=device, dtype=precision) + # FPS for temporal scaling fps = 16 - - forward_batch = ForwardBatch( - data_type="dummy", - ) - + + forward_batch = ForwardBatch(data_type="dummy", ) + logger.info("Running inference...") - + with torch.no_grad(): with torch.autocast('cuda', dtype=precision): # FastVideo model with set_forward_context( - current_timestep=500, - attn_metadata=None, - forward_batch=forward_batch, + current_timestep=500, + attn_metadata=None, + forward_batch=forward_batch, ): # FastVideo expects padding_mask in [B, 1, H, W] format padding_mask_fv = padding_mask.unsqueeze(1) # Add channel dimension for FastVideo # FastVideo supports both [B] and [B, T] formats - use [B, T] to match official model # This ensures each frame gets its own timestep embedding (even if values are the same) timestep_fv = timestep # Already in [B, T] format - + output_fv = fastvideo_model( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, @@ -389,24 +355,24 @@ def test_cosmos25_transformer(): padding_mask=padding_mask_fv, fps=fps, ) - + # Reference model (if available) if reference_model is not None: # Prepare input for reference model # MinimalV1LVGDiT adds condition mask internally, so pass them separately from cosmos_predict2._src.predict2.conditioner import DataType - + # Determine data_type based on temporal dimension num_frames = hidden_states.shape[2] ref_data_type = DataType.VIDEO if num_frames > 1 else DataType.IMAGE - + # Reference model expects different input format # Pass hidden_states without condition_mask (model concatenates it internally) # timestep is already in [B, T] format with correct dtype # padding_mask is already in [B, H, W] format (no channel dimension) # FPS should be a tensor [B] or scalar fps_tensor = torch.tensor([fps], device=device, dtype=precision) - + output_ref = reference_model( x_B_C_T_H_W=hidden_states, # [B, 16, T, H, W] - model will add condition mask timesteps_B_T=timestep, # Already in [B, T] format @@ -416,38 +382,37 @@ def test_cosmos25_transformer(): padding_mask=padding_mask, # [B, H, W] format data_type=ref_data_type, ) - + # Check FastVideo output shape and dtype logger.info(f"FastVideo output shape: {output_fv.shape}") logger.info(f"FastVideo output dtype: {output_fv.dtype}") assert output_fv.shape[0] == batch_size, "Batch size mismatch" assert output_fv.shape[1] == 16, "Output channels should be 16" assert output_fv.dtype == precision, f"Output dtype mismatch: {output_fv.dtype} vs {precision}" - + # Compare with reference if available if reference_model is not None: logger.info(f"Reference output shape: {output_ref.shape}") - + # Check if outputs have the same shape assert output_fv.shape == output_ref.shape, \ f"Output shapes don't match: {output_fv.shape} vs {output_ref.shape}" assert output_fv.dtype == output_ref.dtype, \ f"Output dtype don't match: {output_fv.dtype} vs {output_ref.dtype}" - + # Check if outputs are similar max_diff = torch.max(torch.abs(output_fv - output_ref)) mean_diff = torch.mean(torch.abs(output_fv - output_ref)) relative_diff = mean_diff / (torch.mean(torch.abs(output_ref)) + 1e-8) - + logger.info(f"Max difference: {max_diff.item():.6f}") logger.info(f"Mean difference: {mean_diff.item():.6f}") logger.info(f"Relative difference: {relative_diff.item():.6f}") - - + # Allow for some numerical differences due to implementation details assert max_diff < 1e-1, f"Maximum difference too large: {max_diff.item()}" assert mean_diff < 1e-2, f"Mean difference too large: {mean_diff.item()}" - + logger.info("✓ COSMOS 2.5 FastVideo implementation matches reference!") else: logger.warning("Reference model not available, skipping comparison") @@ -459,10 +424,10 @@ def test_cosmos25_transformer_video(): """Test COSMOS 2.5 transformer with video input (multiple frames).""" device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") precision = torch.bfloat16 - + # Create COSMOS 2.5 specific config from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25ArchConfig - + arch_config = Cosmos25ArchConfig( num_attention_heads=16, attention_head_dim=128, @@ -482,13 +447,13 @@ def test_cosmos25_transformer_video(): rope_enable_fps_modulation=False, qk_norm="rms_norm", ) - + cosmos25_config = Cosmos25VideoConfig(arch_config=arch_config) - + # Create FastVideo model directly (Cosmos 2.5 is not in diffusers format) logger.info("Creating FastVideo COSMOS 2.5 model for video test...") from fastvideo.models.dits.cosmos2_5 import Cosmos25Transformer3DModel - + # Get hf_config from the arch_config for model initialization hf_config = { 'in_channels': arch_config.in_channels, @@ -509,16 +474,16 @@ def test_cosmos25_transformer_video(): 'rope_enable_fps_modulation': arch_config.rope_enable_fps_modulation, 'qk_norm': arch_config.qk_norm, } - + model = Cosmos25Transformer3DModel(config=cosmos25_config, hf_config=hf_config) model = model.to(device, dtype=precision) model.eval() - + # Create video input with multiple frames batch_size = 1 num_frames = 16 # Video with 16 frames seq_len = 77 - + hidden_states = torch.randn( batch_size, 16, @@ -526,54 +491,35 @@ def test_cosmos25_transformer_video(): 64, 64, device=device, - dtype=precision - ) - - condition_mask = torch.zeros( - batch_size, - 1, - num_frames, - 64, - 64, - device=device, - dtype=precision - ) + dtype=precision) + + condition_mask = torch.zeros(batch_size, 1, num_frames, 64, 64, device=device, dtype=precision) # Set first 2 frames as conditioning condition_mask[:, :, :2, :, :] = 1.0 - + encoder_hidden_states = torch.randn( batch_size, seq_len, 100352, # Qwen 7B embedding dimension (matches crossattn_proj input) device=device, - dtype=precision - ) - + dtype=precision) + timestep = torch.tensor([500], device=device, dtype=torch.long) - - padding_mask = torch.ones( - batch_size, - 1, - 64, - 64, - device=device, - dtype=precision - ) - + + padding_mask = torch.ones(batch_size, 1, 64, 64, device=device, dtype=precision) + fps = 16 - - forward_batch = ForwardBatch( - data_type="dummy", - ) - + + forward_batch = ForwardBatch(data_type="dummy", ) + logger.info("Running video inference...") - + with torch.no_grad(): with torch.autocast('cuda', dtype=precision): with set_forward_context( - current_timestep=500, - attn_metadata=None, - forward_batch=forward_batch, + current_timestep=500, + attn_metadata=None, + forward_batch=forward_batch, ): output = model( hidden_states=hidden_states, @@ -583,16 +529,16 @@ def test_cosmos25_transformer_video(): padding_mask=padding_mask, fps=fps, ) - + logger.info(f"Video output shape: {output.shape}") logger.info(f"Video output dtype: {output.dtype}") - + # Check output shape assert output.shape[0] == batch_size, "Batch size mismatch" assert output.shape[1] == 16, "Output channels should be 16" assert output.shape[2] == num_frames, "Number of frames mismatch" assert output.dtype == precision, f"Output dtype mismatch" - + logger.info("✓ COSMOS 2.5 video inference successful!") diff --git a/fastvideo/tests/transformers/test_flux.py b/fastvideo/tests/transformers/test_flux.py index fb6b2797b..902cb6599 100644 --- a/fastvideo/tests/transformers/test_flux.py +++ b/fastvideo/tests/transformers/test_flux.py @@ -60,10 +60,8 @@ requires_cuda = pytest.mark.skipif( requires_weights = pytest.mark.skipif( not glob.glob(os.path.join(_flux_transformer_path(), "*.safetensors")), - reason=( - f"No safetensors under {_flux_transformer_path()} — download FLUX.1-dev " - "transformer or set FLUX_TRANSFORMER_PATH" - ), + reason=(f"No safetensors under {_flux_transformer_path()} — download FLUX.1-dev " + "transformer or set FLUX_TRANSFORMER_PATH"), ) @@ -118,7 +116,7 @@ def test_flux_transformer_parity_vs_diffusers() -> None: # Diffusers pipeline passes scheduler timesteps / 1000 (float, same dtype as latents). timestep = torch.tensor([512.0], device=device, dtype=precision) / 1000.0 - guidance = torch.full((batch_size,), 3.5, device=device, dtype=torch.float32) + guidance = torch.full((batch_size, ), 3.5, device=device, dtype=torch.float32) txt_ids = torch.zeros(text_len, 3, device=device, dtype=torch.long) img_ids = _prepare_latent_image_ids(latent_h, latent_w, device, dtype=torch.long) @@ -130,13 +128,13 @@ def test_flux_transformer_parity_vs_diffusers() -> None: fv_model = loader.load(transformer_path, args).to(device=device, dtype=precision) fv_model.eval() with ( - torch.no_grad(), - torch.amp.autocast("cuda", dtype=precision), - set_forward_context( - current_timestep=512, - attn_metadata=None, - forward_batch=forward_batch, - ), + torch.no_grad(), + torch.amp.autocast("cuda", dtype=precision), + set_forward_context( + current_timestep=512, + attn_metadata=None, + forward_batch=forward_batch, + ), ): fv_out = fv_model( hidden_states=hidden_states.clone(), @@ -153,14 +151,10 @@ def test_flux_transformer_parity_vs_diffusers() -> None: if torch.cuda.is_available(): torch.cuda.empty_cache() - hf_model = ( - HFFluxTransformer2DModel.from_pretrained( - transformer_path, - torch_dtype=precision, - ) - .to(device) - .eval() - ) + hf_model = (HFFluxTransformer2DModel.from_pretrained( + transformer_path, + torch_dtype=precision, + ).to(device).eval()) with torch.no_grad(), torch.amp.autocast("cuda", dtype=precision): hf_out = hf_model( hidden_states=hidden_states.clone(), diff --git a/fastvideo/tests/transformers/test_hunyuangamecraft.py b/fastvideo/tests/transformers/test_hunyuangamecraft.py index db1e5881f..e4e88f461 100644 --- a/fastvideo/tests/transformers/test_hunyuangamecraft.py +++ b/fastvideo/tests/transformers/test_hunyuangamecraft.py @@ -41,10 +41,10 @@ REFERENCE_LATENT = 42351.12903189659 @pytest.mark.usefixtures("distributed_setup") def test_hunyuangamecraft_transformer(): """Test HunyuanGameCraft transformer regression.""" - + if not os.path.exists(TRANSFORMER_PATH): pytest.skip(f"Weights not found at {TRANSFORMER_PATH}") - + sp_rank = get_sp_parallel_rank() sp_world_size = get_sp_world_size() @@ -56,9 +56,7 @@ def test_hunyuangamecraft_transformer(): model_path=TRANSFORMER_PATH, dit_cpu_offload=False, use_fsdp_inference=False, - pipeline_config=PipelineConfig( - dit_config=HunyuanGameCraftConfig(), dit_precision=precision_str - ), + pipeline_config=PipelineConfig(dit_config=HunyuanGameCraftConfig(), dit_precision=precision_str), ) args.device = device @@ -82,10 +80,13 @@ def test_hunyuangamecraft_transformer(): text_seq_len = 32 # Input latents [B, 33, T, H, W] - 16 latent + 16 gt_latent + 1 mask - hidden_states = torch.randn( - batch_size, 33, latent_frames, latent_height, latent_width, - device=device, dtype=precision - ) + hidden_states = torch.randn(batch_size, + 33, + latent_frames, + latent_height, + latent_width, + device=device, + dtype=precision) if sp_world_size > 1: chunk_per_rank = hidden_states.shape[2] // sp_world_size @@ -93,10 +94,10 @@ def test_hunyuangamecraft_transformer(): # Text embeddings (LLaMA) text_states = torch.randn(batch_size, text_seq_len, 4096, device=device, dtype=precision) - + # CLIP pooled embeddings text_states_2 = torch.randn(batch_size, 768, device=device, dtype=precision) - + # Text mask text_mask = torch.ones(batch_size, text_seq_len, device=device, dtype=torch.long) @@ -105,10 +106,7 @@ def test_hunyuangamecraft_transformer(): # Camera states [B, num_frames, 6, H, W] - 33 video frames at full resolution # For 9 latent frames, we use 33 camera frames (matching official model) - camera_states = torch.randn( - batch_size, 33, 6, 352, 640, - device=device, dtype=precision - ) + camera_states = torch.randn(batch_size, 33, 6, 352, 640, device=device, dtype=precision) encoder_hidden_states = [text_states, text_states_2] diff --git a/fastvideo/tests/transformers/test_hyworld.py b/fastvideo/tests/transformers/test_hyworld.py index 91828f476..6e274a595 100644 --- a/fastvideo/tests/transformers/test_hyworld.py +++ b/fastvideo/tests/transformers/test_hyworld.py @@ -55,9 +55,7 @@ def test_hyworld_transformer(): model_path=transformer_path, dit_cpu_offload=False, use_fsdp_inference=False, - pipeline_config=PipelineConfig( - dit_config=HYWorldConfig(), dit_precision=precision_str - ), + pipeline_config=PipelineConfig(dit_config=HYWorldConfig(), dit_precision=precision_str), ) args.device = device @@ -86,57 +84,42 @@ def test_hyworld_transformer(): latent_width = 30 # Video latents [B, C, T, H, W] - hidden_states = torch.randn( - batch_size, 65, num_frames, latent_height, latent_width, - device=device, dtype=precision - ) + hidden_states = torch.randn(batch_size, 65, num_frames, latent_height, latent_width, device=device, dtype=precision) if sp_world_size > 1: chunk_per_rank = hidden_states.shape[2] // sp_world_size hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) * chunk_per_rank] # Text embeddings [B, L, D] (including global token) - encoder_hidden_states = torch.randn( - batch_size, seq_len + 1, 3584, device=device, dtype=precision - ) + encoder_hidden_states = torch.randn(batch_size, seq_len + 1, 3584, device=device, dtype=precision) # Create attention mask for encoder_hidden_states - encoder_attention_mask = torch.ones( - batch_size, seq_len + 1, device=device, dtype=torch.bool - ) + encoder_attention_mask = torch.ones(batch_size, seq_len + 1, device=device, dtype=torch.bool) encoder_hidden_states[:, 15:] = 0 encoder_attention_mask[:, 15:] = False - encoder_hidden_states_2 = torch.randn( - batch_size, seq_len_2 + 1, 1472, device=device, dtype=precision - ) - encoder_attention_mask_2 = torch.ones( - batch_size, seq_len_2 + 1, device=device, dtype=torch.bool - ) + encoder_hidden_states_2 = torch.randn(batch_size, seq_len_2 + 1, 1472, device=device, dtype=precision) + encoder_attention_mask_2 = torch.ones(batch_size, seq_len_2 + 1, device=device, dtype=torch.bool) encoder_hidden_states_2[:, 39:] = 0 encoder_attention_mask_2[:, 39:] = False # Image embeddings - image_embeds = torch.zeros( - batch_size, 729, 1152, dtype=precision, device=device - ) + image_embeds = torch.zeros(batch_size, 729, 1152, dtype=precision, device=device) # Action tensor [B*T] - discrete action per frame, first frame of each batch is 0 - action = torch.randint(0, 10, (batch_size * num_frames,), device=device) + action = torch.randint(0, 10, (batch_size * num_frames, ), device=device) action[::num_frames] = 0 # First frame of each batch is 0 action = action.to(dtype=precision) # Camera view matrices [B, T, 4, 4] - 4x4 extrinsic/view matrices - viewmats = torch.eye(4, dtype=precision, device=device).unsqueeze(0).unsqueeze(0).expand( - batch_size, num_frames, -1, -1 - ).contiguous() + viewmats = torch.eye(4, dtype=precision, + device=device).unsqueeze(0).unsqueeze(0).expand(batch_size, num_frames, -1, -1).contiguous() # Camera intrinsics [B, T, 3, 3] - 3x3 intrinsic matrices - Ks = torch.eye(3, dtype=precision, device=device).unsqueeze(0).unsqueeze(0).expand( - batch_size, num_frames, -1, -1 - ).contiguous() + Ks = torch.eye(3, dtype=precision, device=device).unsqueeze(0).unsqueeze(0).expand(batch_size, num_frames, -1, + -1).contiguous() # Timestep [B*T] - one timestep value per frame - timestep = torch.full((batch_size * num_frames,), 500, device=device, dtype=precision) + timestep = torch.full((batch_size * num_frames, ), 500, device=device, dtype=precision) # Timestep for text [B] - one timestep value per batch timestep_txt = torch.tensor([500], device=device, dtype=precision) @@ -145,9 +128,9 @@ def test_hyworld_transformer(): with torch.no_grad(): with torch.amp.autocast("cuda", dtype=precision): with set_forward_context( - current_timestep=0, - attn_metadata=None, - forward_batch=forward_batch, + current_timestep=0, + attn_metadata=None, + forward_batch=forward_batch, ): output = model( hidden_states=hidden_states, diff --git a/fastvideo/tests/transformers/test_wanvideo.py b/fastvideo/tests/transformers/test_wanvideo.py index 9f93783f7..6a6e2f509 100644 --- a/fastvideo/tests/transformers/test_wanvideo.py +++ b/fastvideo/tests/transformers/test_wanvideo.py @@ -16,7 +16,6 @@ from fastvideo.utils import maybe_download_model from fastvideo.configs.models.dits import WanVideoConfig from fastvideo.pipelines.pipeline_batch_info import ForwardBatch - logger = init_logger(__name__) os.environ["MASTER_ADDR"] = "localhost" @@ -42,11 +41,9 @@ def test_wan_transformer(): loader = TransformerLoader() model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision) - model1 = ( - WanTransformer3DModel.from_pretrained(TRANSFORMER_PATH, device=device, torch_dtype=precision) - .to(device, dtype=precision) - .requires_grad_(False) - ) + model1 = (WanTransformer3DModel.from_pretrained(TRANSFORMER_PATH, device=device, + torch_dtype=precision).to(device, + dtype=precision).requires_grad_(False)) total_params = sum(p.numel() for p in model1.parameters()) # Calculate weight sum for model1 (converting to float64 to avoid overflow) @@ -86,9 +83,7 @@ def test_wan_transformer(): # Timestep timestep = torch.tensor([500], device=device, dtype=precision) - forward_batch = ForwardBatch( - data_type="dummy", - ) + forward_batch = ForwardBatch(data_type="dummy", ) with torch.amp.autocast("cuda", dtype=precision): output1 = model1( @@ -98,13 +93,13 @@ def test_wan_transformer(): return_dict=False, )[0] with set_forward_context( - current_timestep=0, - attn_metadata=None, - forward_batch=forward_batch, + current_timestep=0, + attn_metadata=None, + forward_batch=forward_batch, ): - output2 = model2( - hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, timestep=timestep - ) + output2 = model2(hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep) # Check if outputs have the same shape assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" diff --git a/fastvideo/tests/utils.py b/fastvideo/tests/utils.py index 6ae84a273..d80c2198f 100644 --- a/fastvideo/tests/utils.py +++ b/fastvideo/tests/utils.py @@ -132,8 +132,8 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True): # Process each frame individually for i in range(min_frames): - img1 = frames1[i : i + 1] - img2 = frames2[i : i + 1] + img1 = frames1[i:i + 1] + img2 = frames2[i:i + 1] with torch.no_grad(): value = ms_ssim(img1, img2, data_range=1.0) if use_ms_ssim else ssim(img1, img2, data_range=1.0) @@ -220,7 +220,10 @@ def write_ssim_results(output_dir, ssim_values, reference_path, generated_path, "max_ssim": max_ssim, "reference_video": reference_path, "generated_video": generated_path, - "parameters": {"num_inference_steps": num_inference_steps, "prompt": prompt}, + "parameters": { + "num_inference_steps": num_inference_steps, + "prompt": prompt + }, } test_name = f"steps{num_inference_steps}_{prompt[:100]}" diff --git a/fastvideo/tests/vaes/test_hunyuan15_vae.py b/fastvideo/tests/vaes/test_hunyuan15_vae.py index 731ebff45..ab1c0ddf2 100644 --- a/fastvideo/tests/vaes/test_hunyuan15_vae.py +++ b/fastvideo/tests/vaes/test_hunyuan15_vae.py @@ -22,9 +22,7 @@ os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29503" BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" -MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, - local_dir=os.path.join( - "data", BASE_MODEL_PATH)) +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH)) VAE_PATH = os.path.join(MODEL_PATH, "vae") CONFIG_PATH = os.path.join(VAE_PATH, "config.json") @@ -34,12 +32,12 @@ def test_hunyuan_vae(): device = torch.device("cuda:0") precision = torch.float32 precision_str = "fp32" - args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=Hunyuan15VAEConfig(), vae_precision=precision_str)) + args = FastVideoArgs(model_path=VAE_PATH, + pipeline_config=PipelineConfig(vae_config=Hunyuan15VAEConfig(), vae_precision=precision_str)) args.device = device args.vae_cpu_offload = False - model1 = AutoencoderKLHunyuanVideo15.from_pretrained( - VAE_PATH, torch_dtype=precision).to(device).eval() + model1 = AutoencoderKLHunyuanVideo15.from_pretrained(VAE_PATH, torch_dtype=precision).to(device).eval() model1.enable_tiling() loader = VAELoader() @@ -50,13 +48,7 @@ def test_hunyuan_vae(): batch_size = 1 # Video input [B, C, T, H, W] - input_tensor = torch.randn(batch_size, - 3, - 81, - 512, - 512, - device=device, - dtype=precision) + input_tensor = torch.randn(batch_size, 3, 81, 512, 512, device=device, dtype=precision) # Disable gradients for inference with torch.no_grad(): @@ -66,10 +58,8 @@ def test_hunyuan_vae(): assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}" max_diff_encode = torch.max(torch.abs(latent1.float() - latent2.float())) mean_diff_encode = torch.mean(torch.abs(latent1.float() - latent2.float())) - logger.info("Maximum difference between encoded latents: %s", - max_diff_encode.item()) - logger.info("Mean difference between encoded latents: %s", - mean_diff_encode.item()) + logger.info("Maximum difference between encoded latents: %s", max_diff_encode.item()) + logger.info("Mean difference between encoded latents: %s", mean_diff_encode.item()) assert max_diff_encode < 1e-5, f"Encoded latents differ significantly: max diff = {max_diff_encode.item()}, mean diff = {mean_diff_encode.item()}" # Test decoding @@ -83,9 +73,6 @@ def test_hunyuan_vae(): assert video1.shape == video2.shape, f"Video shapes don't match: {video1.shape} vs {video2.shape}" max_diff_decode = torch.max(torch.abs(video1.float() - video2.float())) mean_diff_decode = torch.mean(torch.abs(video1.float() - video2.float())) - logger.info("Maximum difference between decoded videos: %s", - max_diff_decode.item()) - logger.info("Mean difference between decoded videos: %s", - mean_diff_decode.item()) + logger.info("Maximum difference between decoded videos: %s", max_diff_decode.item()) + logger.info("Mean difference between decoded videos: %s", mean_diff_decode.item()) assert max_diff_decode < 1e-5, f"Decoded videos differ significantly: max diff = {max_diff_decode.item()}, mean diff = {mean_diff_decode.item()}" - diff --git a/fastvideo/tests/vaes/test_wan_vae.py b/fastvideo/tests/vaes/test_wan_vae.py index 31a1175b8..b60973ac1 100644 --- a/fastvideo/tests/vaes/test_wan_vae.py +++ b/fastvideo/tests/vaes/test_wan_vae.py @@ -29,9 +29,8 @@ def test_wan_vae(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") precision = torch.float32 precision_str = "fp32" - args = FastVideoArgs( - model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str) - ) + args = FastVideoArgs(model_path=VAE_PATH, + pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str)) args.device = device args.vae_cpu_offload = False @@ -63,21 +62,16 @@ def test_wan_vae(): latent2 = model2.encode(input_tensor) # Check if latents have the same shape assert latent1.mean.shape == latent2.mean.shape, ( - f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}" - ) + f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}") # Check if latents are similar assert_close(latent1.mean, latent2.mean, atol=1e-4, rtol=1e-4) # Test decoding logger.info("Testing decoding...") latent1_tensor = latent1.mode() - mean1 = ( - torch.tensor(model1.config.latents_mean) - .view(1, model1.config.z_dim, 1, 1, 1) - .to(input_tensor.device, input_tensor.dtype) - ) + mean1 = (torch.tensor(model1.config.latents_mean).view(1, model1.config.z_dim, 1, 1, + 1).to(input_tensor.device, input_tensor.dtype)) std1 = (1.0 / torch.tensor(model1.config.latents_std).view(1, model1.config.z_dim, 1, 1, 1)).to( - input_tensor.device, input_tensor.dtype - ) + input_tensor.device, input_tensor.dtype) latent1_tensor = latent1_tensor / std1 + mean1 output1 = model1.decode(latent1_tensor).sample diff --git a/fastvideo/tests/vaes/test_wan_vae_compile.py b/fastvideo/tests/vaes/test_wan_vae_compile.py index 0c87eba58..ede4dcbb6 100644 --- a/fastvideo/tests/vaes/test_wan_vae_compile.py +++ b/fastvideo/tests/vaes/test_wan_vae_compile.py @@ -12,8 +12,7 @@ encoder/decoder submodules by name and type. """ import torch.nn as nn -from fastvideo.models.vaes.wanvae import (AutoencoderKLWan, WanDecoder3d, - WanEncoder3d, _is_wan_vae_codec) +from fastvideo.models.vaes.wanvae import (AutoencoderKLWan, WanDecoder3d, WanEncoder3d, _is_wan_vae_codec) def test_wan_vae_registers_compile_conditions(): diff --git a/fastvideo/tests/workflow/test_kandinsky5_overfit_caption.py b/fastvideo/tests/workflow/test_kandinsky5_overfit_caption.py index c8ae97f3f..76c61364b 100644 --- a/fastvideo/tests/workflow/test_kandinsky5_overfit_caption.py +++ b/fastvideo/tests/workflow/test_kandinsky5_overfit_caption.py @@ -38,7 +38,9 @@ def test_get_caption_uses_first_list_element(): [""], None, 123, - {"nested": "dict"}, + { + "nested": "dict" + }, ]) def test_get_caption_rejects_invalid_values(bad_cap): with pytest.raises(ValueError): diff --git a/fastvideo/tests/workflow/test_minimax_h3_overfit_preprocess.py b/fastvideo/tests/workflow/test_minimax_h3_overfit_preprocess.py index b2bb9bc9c..8fe8fb51e 100644 --- a/fastvideo/tests/workflow/test_minimax_h3_overfit_preprocess.py +++ b/fastvideo/tests/workflow/test_minimax_h3_overfit_preprocess.py @@ -18,11 +18,8 @@ from fastvideo.pipelines.preprocess.preprocess_minimax_h3_overfit import ( write_parquet, ) - -_TRAINING_CAPTION = ( - "A watermelon wearing a helmet is crushed by a hydraulic press, " - "causing it to flatten and burst open." -) +_TRAINING_CAPTION = ("A watermelon wearing a helmet is crushed by a hydraulic press, " + "causing it to flatten and burst open.") def test_load_crush_smol_training_sample_selects_fixed_record(tmp_path: Path) -> None: @@ -32,16 +29,17 @@ def test_load_crush_smol_training_sample_selects_fixed_record(tmp_path: Path) -> video_path = video_dir / TRAINING_VIDEO_NAME video_path.touch() manifest_path = tmp_path / "videos2caption.json" - manifest_path.write_text(json.dumps([ - { - "path": "another-video.mp4", - "cap": ["Another caption."], - }, - { - "path": TRAINING_VIDEO_NAME, - "cap": [_TRAINING_CAPTION], - }, - ])) + manifest_path.write_text( + json.dumps([ + { + "path": "another-video.mp4", + "cap": ["Another caption."], + }, + { + "path": TRAINING_VIDEO_NAME, + "cap": [_TRAINING_CAPTION], + }, + ])) selected_video_path, caption = load_crush_smol_training_sample(manifest_path, video_dir) diff --git a/fastvideo/third_party/eval/glmasr/modeling_glmasr.py b/fastvideo/third_party/eval/glmasr/modeling_glmasr.py index 9b208b663..706688503 100644 --- a/fastvideo/third_party/eval/glmasr/modeling_glmasr.py +++ b/fastvideo/third_party/eval/glmasr/modeling_glmasr.py @@ -42,7 +42,6 @@ from ._compat import ( ) from .configuration_glmasr import GlmAsrConfig, GlmAsrEncoderConfig - if is_torch_available(): import torch from torch import nn @@ -94,9 +93,7 @@ class GlmAsrRotaryEmbedding(nn.Module): attention_factor = 1.0 # Unused in this type of RoPE # Compute the inverse frequencies - inv_freq = 1.0 / ( - base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) - ) + inv_freq = 1.0 / (base**(torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)) return inv_freq, attention_factor @torch.no_grad() @@ -117,8 +114,8 @@ class GlmAsrRotaryEmbedding(nn.Module): def rotate_half(x): """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] + x1 = x[..., :x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) @@ -211,9 +208,8 @@ class GlmAsrAttention(nn.Module): cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( - self.config._attn_implementation, eager_attention_forward - ) + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(self.config._attn_implementation, + eager_attention_forward) attn_output, attn_weights = attention_interface( self, @@ -232,6 +228,7 @@ class GlmAsrAttention(nn.Module): class GlmAsrMLP(nn.Module): + def __init__(self, config): super().__init__() self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) @@ -246,6 +243,7 @@ class GlmAsrMLP(nn.Module): class GlmAsrEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: GlmAsrConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size @@ -309,8 +307,7 @@ class GlmAsrEncoder(GlmAsrPreTrainedModel): self.conv2 = nn.Conv1d(config.hidden_size, config.hidden_size, kernel_size=3, stride=2, padding=1) self.layers = nn.ModuleList( - [GlmAsrEncoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] - ) + [GlmAsrEncoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) self.norm = nn.LayerNorm(config.hidden_size) self.rotary_emb = GlmAsrRotaryEmbedding(config=config) self.gradient_checkpointing = False @@ -325,9 +322,9 @@ class GlmAsrEncoder(GlmAsrPreTrainedModel): inputs_embeds = inputs_embeds.transpose(1, 2) hidden_states = inputs_embeds - position_embeddings = self.rotary_emb( - hidden_states, position_ids=torch.arange(hidden_states.shape[1], device=hidden_states.device)[None, :] - ) + position_embeddings = self.rotary_emb(hidden_states, + position_ids=torch.arange(hidden_states.shape[1], + device=hidden_states.device)[None, :]) for encoder_layer in self.layers: hidden_states = encoder_layer(hidden_states, position_embeddings=position_embeddings, **kwargs) @@ -355,11 +352,9 @@ class GlmAsrMultiModalProjector(nn.Module): return hidden_states -@auto_docstring( - custom_intro=""" +@auto_docstring(custom_intro=""" The GlmAsr model which consists of a fine-tuned Whisper encoder, a multi-modal projector and a Llama language model. - """ -) + """) class GlmAsrForConditionalGeneration(GlmAsrPreTrainedModel, GenerationMixin): _keep_in_fp32_modules_strict = None _supports_attention_backend = True @@ -396,8 +391,8 @@ class GlmAsrForConditionalGeneration(GlmAsrPreTrainedModel, GenerationMixin): @can_return_tuple @auto_docstring( - custom_intro="Compute audio embeddings from log-mel input features using the audio encoder and multi-modal projector." - ) + custom_intro= + "Compute audio embeddings from log-mel input features using the audio encoder and multi-modal projector.") def get_audio_features( self, input_features: torch.FloatTensor, @@ -416,9 +411,8 @@ class GlmAsrForConditionalGeneration(GlmAsrPreTrainedModel, GenerationMixin): """ audio_outputs = self.audio_tower(input_features, return_dict=True, **kwargs) audio_hidden_states = audio_outputs.last_hidden_state - audio_hidden_states = audio_hidden_states.reshape( - input_features.shape[0], -1, self.config.audio_config.intermediate_size - ) + audio_hidden_states = audio_hidden_states.reshape(input_features.shape[0], -1, + self.config.audio_config.intermediate_size) audio_embeds = self.multi_modal_projector(audio_hidden_states) audio_lengths = input_features_mask.sum(-1) @@ -485,9 +479,8 @@ class GlmAsrForConditionalGeneration(GlmAsrPreTrainedModel, GenerationMixin): # replace text-audio token placeholders with audio embeddings audio_token_mask = (input_ids == self.config.audio_token_id).unsqueeze(-1) - inputs_embeds = inputs_embeds.masked_scatter( - audio_token_mask.to(inputs_embeds.device), audio_embeds.to(inputs_embeds.device) - ) + inputs_embeds = inputs_embeds.masked_scatter(audio_token_mask.to(inputs_embeds.device), + audio_embeds.to(inputs_embeds.device)) outputs: CausalLMOutputWithPast = self.language_model( inputs_embeds=inputs_embeds, diff --git a/fastvideo/third_party/eval/glmasr/processing_glmasr.py b/fastvideo/third_party/eval/glmasr/processing_glmasr.py index b10406402..6ebf69f9a 100644 --- a/fastvideo/third_party/eval/glmasr/processing_glmasr.py +++ b/fastvideo/third_party/eval/glmasr/processing_glmasr.py @@ -18,7 +18,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - import re import numpy as np @@ -29,11 +28,9 @@ from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpa from transformers.tokenization_utils_base import TextInput from transformers.utils import is_torch_available, logging - if is_torch_available(): import torch - logger = logging.get_logger(__name__) @@ -63,7 +60,6 @@ class GlmAsrProcessor(ProcessorMixin): attributes = ["feature_extractor", "tokenizer"] feature_extractor_class = "WhisperFeatureExtractor" tokenizer_class = "PreTrainedTokenizerFast" - r""" Constructs an GlmAsr processor which wraps an GlmAsr feature extractor and an GlmAsr tokenizer into a single processor. @@ -319,8 +315,7 @@ class GlmAsrProcessor(ProcessorMixin): elif isinstance(prompt, (list, tuple)): if len(prompt) != batch_size: raise ValueError( - f"Received {len(prompt)} prompt(s) for {batch_size} audio sample(s); counts must match." - ) + f"Received {len(prompt)} prompt(s) for {batch_size} audio sample(s); counts must match.") prompts = [] for item in prompt: if item is None: @@ -332,20 +327,23 @@ class GlmAsrProcessor(ProcessorMixin): else: raise TypeError("`prompt` must be a string, a sequence of strings, or `None`.") - conversations = [ - [ + conversations = [[{ + "role": + "user", + "content": [ { - "role": "user", - "content": [ - {"type": "audio", "path": audio_item} - if isinstance(audio_item, str) - else {"type": "audio", "audio": audio_item}, - {"type": "text", "text": prompt_text}, - ], - } - ] - for prompt_text, audio_item in zip(prompts, audio_items) - ] + "type": "audio", + "path": audio_item + } if isinstance(audio_item, str) else { + "type": "audio", + "audio": audio_item + }, + { + "type": "text", + "text": prompt_text + }, + ], + }] for prompt_text, audio_item in zip(prompts, audio_items)] return self.apply_chat_template( conversations, @@ -382,12 +380,12 @@ class GlmAsrProcessor(ProcessorMixin): stripped = text.strip() for prefix in ( - "The spoken content of the audio is", - "The transcription of the audio is", - "The content of the input audio is", + "The spoken content of the audio is", + "The transcription of the audio is", + "The content of the input audio is", ): if stripped.startswith(prefix): - stripped = stripped[len(prefix) :].strip() + stripped = stripped[len(prefix):].strip() break if stripped.endswith("."): diff --git a/fastvideo/third_party/eval/synchformer/ast.py b/fastvideo/third_party/eval/synchformer/ast.py index 981a6d66f..7c15d9ad6 100644 --- a/fastvideo/third_party/eval/synchformer/ast.py +++ b/fastvideo/third_party/eval/synchformer/ast.py @@ -5,7 +5,7 @@ from transformers.modeling_outputs import BaseModelOutputWithPooling # importing modified version of AST from fastvideo.third_party.eval.synchformer.hf_src.modeling_ast import (ASTConfig, ASTForAudioClassification) from fastvideo.third_party.eval.synchformer.motionformer import (AveragePooling, BaseEncoderLayer, - TemporalTransformerEncoderLayer) + TemporalTransformerEncoderLayer) class AST(torch.nn.Module): @@ -90,11 +90,10 @@ class AST(torch.nn.Module): # we can reuse the same layer as for temporal factorization (B, dim_to_agg, D) -> (B, D) # we need to add pos emb (PE) because previously we added the same PE for each segment pos_max_len = max_segments if max_segments is not None else 16 # 16 = 10sec//0.64sec + 1 - self.global_attn_agg = TemporalTransformerEncoderLayer( - add_pos_emb=True, - pos_emb_drop=self.config.hidden_dropout_prob, - pos_max_len=pos_max_len, - **transf_enc_layer_kwargs) + self.global_attn_agg = TemporalTransformerEncoderLayer(add_pos_emb=True, + pos_emb_drop=self.config.hidden_dropout_prob, + pos_max_len=pos_max_len, + **transf_enc_layer_kwargs) elif agg_segments_module == 'AveragePooling': self.global_attn_agg = AveragePooling(avg_pattern='B S D -> B D') else: @@ -132,10 +131,7 @@ class AST(torch.nn.Module): orig_shape_s = (B, 1, T, F) # NOTE: since x is (B, S, T, F), and forward_segments expects (BS, T, F). # (B, S, T, F)[:, s] is (B, T, F) or (BS, T, F) if S=1. - x = torch.cat([ - self.forward_segments(x[:, s], orig_shape_s, **ast_kwargs).unsqueeze(1) - for s in range(S) - ], + x = torch.cat([self.forward_segments(x[:, s], orig_shape_s, **ast_kwargs).unsqueeze(1) for s in range(S)], dim=1) else: orig_shape = (B, S, T, F) @@ -168,8 +164,7 @@ class AST(torch.nn.Module): if cont_mask is not None: # duplicating the mask for the latent dimension (D) to be compatible with the next func x_mask = x_mask.unsqueeze(-1).expand(-1, -1, self.config.hidden_size) - x_mask = self.restore_freq_temp_dims(x_mask, - orig_shape) # (BS, D, f, t) <- (B*S, T, D) + x_mask = self.restore_freq_temp_dims(x_mask, orig_shape) # (BS, D, f, t) <- (B*S, T, D) # again removing the latent x_mask = x_mask[:, 0, :, :] else: @@ -222,8 +217,7 @@ class AST(torch.nn.Module): if self.max_spec_t is not None: self.config.max_length = self.max_spec_t f, t = self.ast.embeddings.get_shape(self.config) - shortened = self.ast.embeddings.position_embeddings[:, :f * t + 2].clone( - ) # +2 for CLS and distill tokens + shortened = self.ast.embeddings.position_embeddings[:, :f * t + 2].clone() # +2 for CLS and distill tokens self.ast.embeddings.position_embeddings = torch.nn.Parameter(shortened).to(self.device) def to(self, device): diff --git a/fastvideo/third_party/eval/synchformer/hf_src/modeling_ast.py b/fastvideo/third_party/eval/synchformer/hf_src/modeling_ast.py index 19ec981bf..d363edfa9 100644 --- a/fastvideo/third_party/eval/synchformer/hf_src/modeling_ast.py +++ b/fastvideo/third_party/eval/synchformer/hf_src/modeling_ast.py @@ -14,7 +14,6 @@ # limitations under the License. # Modified by v-iashin to support token masking - """ PyTorch Audio Spectrogram Transformer (AST) model.""" import math @@ -66,7 +65,6 @@ _SEQ_CLASS_CHECKPOINT = "MIT/ast-finetuned-audioset-10-10-0.4593" _SEQ_CLASS_EXPECTED_OUTPUT = "'Speech'" _SEQ_CLASS_EXPECTED_LOSS = 0.17 - AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = [ "MIT/ast-finetuned-audioset-10-10-0.4593", # See all Audio Spectrogram Transformer models at https://huggingface.co/models?filter=ast @@ -125,9 +123,10 @@ class ASTPatchEmbeddings(nn.Module): frequency_stride = config.frequency_stride time_stride = config.time_stride - self.projection = nn.Conv2d( - 1, config.hidden_size, kernel_size=(patch_size, patch_size), stride=(frequency_stride, time_stride) - ) + self.projection = nn.Conv2d(1, + config.hidden_size, + kernel_size=(patch_size, patch_size), + stride=(frequency_stride, time_stride)) def forward(self, input_values: torch.Tensor) -> torch.Tensor: input_values = input_values.unsqueeze(1) @@ -138,13 +137,12 @@ class ASTPatchEmbeddings(nn.Module): # Copied from transformers.models.vit.modeling_vit.ViTSelfAttention with ViT->AST class ASTSelfAttention(nn.Module): + def __init__(self, config: ASTConfig) -> None: super().__init__() if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): - raise ValueError( - f"The hidden size {config.hidden_size,} is not a multiple of the number of attention " - f"heads {config.num_attention_heads}." - ) + raise ValueError(f"The hidden size {config.hidden_size,} is not a multiple of the number of attention " + f"heads {config.num_attention_heads}.") self.num_attention_heads = config.num_attention_heads self.attention_head_size = int(config.hidden_size / config.num_attention_heads) @@ -161,10 +159,11 @@ class ASTSelfAttention(nn.Module): x = x.view(new_x_shape) return x.permute(0, 2, 1, 3) - def forward( - self, hidden_states, tok_mask: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False - ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: + def forward(self, + hidden_states, + tok_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: mixed_query_layer = self.query(hidden_states) key_layer = self.transpose_for_scores(self.key(hidden_states)) @@ -195,10 +194,10 @@ class ASTSelfAttention(nn.Module): context_layer = torch.matmul(attention_probs, value_layer) context_layer = context_layer.permute(0, 2, 1, 3).contiguous() - new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size, ) context_layer = context_layer.view(new_context_layer_shape) - outputs = (context_layer, attention_probs) if output_attentions else (context_layer,) + outputs = (context_layer, attention_probs) if output_attentions else (context_layer, ) return outputs @@ -224,6 +223,7 @@ class ASTSelfOutput(nn.Module): # Copied from transformers.models.vit.modeling_vit.ViTAttention with ViT->AST class ASTAttention(nn.Module): + def __init__(self, config: ASTConfig) -> None: super().__init__() self.attention = ASTSelfAttention(config) @@ -233,9 +233,8 @@ class ASTAttention(nn.Module): def prune_heads(self, heads: Set[int]) -> None: if len(heads) == 0: return - heads, index = find_pruneable_heads_and_indices( - heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads - ) + heads, index = find_pruneable_heads_and_indices(heads, self.attention.num_attention_heads, + self.attention.attention_head_size, self.pruned_heads) # Prune linear layers self.attention.query = prune_linear_layer(self.attention.query, index) @@ -259,12 +258,13 @@ class ASTAttention(nn.Module): attention_output = self.output(self_outputs[0], hidden_states) - outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them + outputs = (attention_output, ) + self_outputs[1:] # add attentions if we output them return outputs # Copied from transformers.models.vit.modeling_vit.ViTIntermediate with ViT->AST class ASTIntermediate(nn.Module): + def __init__(self, config: ASTConfig) -> None: super().__init__() self.dense = nn.Linear(config.hidden_size, config.intermediate_size) @@ -282,6 +282,7 @@ class ASTIntermediate(nn.Module): # Copied from transformers.models.vit.modeling_vit.ViTOutput with ViT->AST class ASTOutput(nn.Module): + def __init__(self, config: ASTConfig) -> None: super().__init__() self.dense = nn.Linear(config.intermediate_size, config.hidden_size) @@ -336,13 +337,14 @@ class ASTLayer(nn.Module): # second residual connection is done here layer_output = self.output(layer_output, hidden_states) - outputs = (layer_output,) + outputs + outputs = (layer_output, ) + outputs return outputs # Copied from transformers.models.vit.modeling_vit.ViTEncoder with ViT->AST class ASTEncoder(nn.Module): + def __init__(self, config: ASTConfig) -> None: super().__init__() self.config = config @@ -363,13 +365,14 @@ class ASTEncoder(nn.Module): for i, layer_module in enumerate(self.layer): if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) + all_hidden_states = all_hidden_states + (hidden_states, ) layer_head_mask = head_mask[i] if head_mask is not None else None if self.gradient_checkpointing and self.training: def create_custom_forward(module): + def custom_forward(*inputs): return module(*inputs, output_attentions) @@ -387,10 +390,10 @@ class ASTEncoder(nn.Module): hidden_states = layer_outputs[0] if output_attentions: - all_self_attentions = all_self_attentions + (layer_outputs[1],) + all_self_attentions = all_self_attentions + (layer_outputs[1], ) if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) + all_hidden_states = all_hidden_states + (hidden_states, ) if not return_dict: return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) @@ -418,9 +421,9 @@ class ASTPreTrainedModel(PreTrainedModel): if isinstance(module, (nn.Linear, nn.Conv2d)): # Upcast the input in `fp32` and cast it back to desired `dtype` to avoid # `trunc_normal_cpu` not implemented in `half` issues - module.weight.data = nn.init.trunc_normal_( - module.weight.data.to(torch.float32), mean=0.0, std=self.config.initializer_range - ).to(module.weight.dtype) + module.weight.data = nn.init.trunc_normal_(module.weight.data.to(torch.float32), + mean=0.0, + std=self.config.initializer_range).to(module.weight.dtype) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.LayerNorm): @@ -473,6 +476,7 @@ AUDIO_SPECTROGRAM_TRANSFORMER_INPUTS_DOCSTRING = r""" AUDIO_SPECTROGRAM_TRANSFORMER_START_DOCSTRING, ) class ASTModel(ASTPreTrainedModel): + def __init__(self, config: ASTConfig): super().__init__(config) self.config = config @@ -514,9 +518,8 @@ class ASTModel(ASTPreTrainedModel): return_dict: Optional[bool] = None, ): output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else self.config.output_hidden_states) return_dict = return_dict if return_dict is not None else self.config.use_return_dict if input_values is None: @@ -551,7 +554,7 @@ class ASTModel(ASTPreTrainedModel): # replace nans with 0s; these are the tokens that correspond to the masked content tok_mask = ~torch.isnan(indicator) # since all values in the D-dimension (latent) will also be nans, we can just use the first el - tok_mask = tok_mask[:, :, 0] # (BS, 2+num_patches) -- 2 is from CLS and DISTIL tokens + tok_mask = tok_mask[:, :, 0] # (BS, 2+num_patches) -- 2 is from CLS and DISTIL tokens else: tok_mask = None @@ -580,11 +583,11 @@ class ASTModel(ASTPreTrainedModel): class ASTMLPHead(nn.Module): + def __init__(self, config: ASTConfig): super().__init__() self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - self.dense = nn.Linear( - config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity() + self.dense = nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity() def forward(self, hidden_state): hidden_state = self.layernorm(hidden_state) @@ -600,6 +603,7 @@ class ASTMLPHead(nn.Module): AUDIO_SPECTROGRAM_TRANSFORMER_START_DOCSTRING, ) class ASTForAudioClassification(ASTPreTrainedModel): + def __init__(self, config: ASTConfig) -> None: super().__init__(config) @@ -675,8 +679,8 @@ class ASTForAudioClassification(ASTPreTrainedModel): loss = loss_fct(logits, labels) if not return_dict: - output = (logits,) + outputs[2:] - return ((loss,) + output) if loss is not None else output + output = (logits, ) + outputs[2:] + return ((loss, ) + output) if loss is not None else output return SequenceClassifierOutput( loss=loss, diff --git a/fastvideo/third_party/eval/synchformer/motionformer.py b/fastvideo/third_party/eval/synchformer/motionformer.py index 1444366b8..1c3985348 100644 --- a/fastvideo/third_party/eval/synchformer/motionformer.py +++ b/fastvideo/third_party/eval/synchformer/motionformer.py @@ -19,12 +19,9 @@ FILE2URL = { 'divided_224_16x4.yaml': 'https://raw.githubusercontent.com/facebookresearch/Motionformer/bf43d50/configs/SSV2/divided_224_16x4.yaml', # ckpt - 'ssv2_motionformer_224_16x4.pyth': - 'https://dl.fbaipublicfiles.com/motionformer/ssv2_motionformer_224_16x4.pyth', - 'ssv2_joint_224_16x4.pyth': - 'https://dl.fbaipublicfiles.com/motionformer/ssv2_joint_224_16x4.pyth', - 'ssv2_divided_224_16x4.pyth': - 'https://dl.fbaipublicfiles.com/motionformer/ssv2_divided_224_16x4.pyth', + 'ssv2_motionformer_224_16x4.pyth': 'https://dl.fbaipublicfiles.com/motionformer/ssv2_motionformer_224_16x4.pyth', + 'ssv2_joint_224_16x4.pyth': 'https://dl.fbaipublicfiles.com/motionformer/ssv2_joint_224_16x4.pyth', + 'ssv2_divided_224_16x4.pyth': 'https://dl.fbaipublicfiles.com/motionformer/ssv2_divided_224_16x4.pyth', } @@ -69,8 +66,7 @@ class MotionFormer(VisionTransformer): # init from motionformer ckpt or from our Stage I ckpt # depending on whether the feat extractor was pre-trained on AVCLIPMoCo or not, we need to # load the state dict differently - was_pt_on_avclip = self.ckpt_path.endswith( - '.pt') # checks if it is a stage I ckpt (FIXME: a bit generic) + was_pt_on_avclip = self.ckpt_path.endswith('.pt') # checks if it is a stage I ckpt (FIXME: a bit generic) if self.ckpt_path.endswith(tuple(mformer_ckpt2cfg.keys())): cfg_fname = mformer_ckpt2cfg[Path(self.ckpt_path).name] elif was_pt_on_avclip: @@ -116,8 +112,7 @@ class MotionFormer(VisionTransformer): # load the ckpt now if ckpt is provided and not from AVCLIPMoCo-pretrained ckpt if (self.ckpt_path is not None) and (not was_pt_on_avclip): _ckpt_load_status = self.load_state_dict(ckpt['model_state'], strict=False) - if len(_ckpt_load_status.missing_keys) > 0 or len( - _ckpt_load_status.unexpected_keys) > 0: + if len(_ckpt_load_status.missing_keys) > 0 or len(_ckpt_load_status.unexpected_keys) > 0: logging.warning(f'Loading exact vfeat_extractor ckpt from {self.ckpt_path} failed.' \ f'Missing keys: {_ckpt_load_status.missing_keys}, ' \ f'Unexpected keys: {_ckpt_load_status.unexpected_keys}') @@ -125,8 +120,7 @@ class MotionFormer(VisionTransformer): logging.info(f'Loading vfeat_extractor ckpt from {self.ckpt_path} succeeded.') if self.extract_features: - assert isinstance(self.norm, - nn.LayerNorm), 'early x[:, 1:, :] may not be safe for per-tr weights' + assert isinstance(self.norm, nn.LayerNorm), 'early x[:, 1:, :] may not be safe for per-tr weights' # pre-logits are Sequential(nn.Linear(emb, emd), act) and `act` is tanh but see the logger self.pre_logits = nn.Identity() # we don't need the classification head (saving memory) @@ -146,8 +140,7 @@ class MotionFormer(VisionTransformer): # define adapters if needed if self.factorize_space_time: if agg_space_module == 'TransformerEncoderLayer': - self.spatial_attn_agg = SpatialTransformerEncoderLayer( - **transf_enc_layer_kwargs) + self.spatial_attn_agg = SpatialTransformerEncoderLayer(**transf_enc_layer_kwargs) elif agg_space_module == 'AveragePooling': self.spatial_attn_agg = AveragePooling(avg_pattern='BS D t h w -> BS D t', then_permute_pattern='BS D t -> BS t D') @@ -164,11 +157,10 @@ class MotionFormer(VisionTransformer): # we can reuse the same layer as for temporal factorization (B, dim_to_agg, D) -> (B, D) # we need to add pos emb (PE) because previously we added the same PE for each segment pos_max_len = max_segments if max_segments is not None else 16 # 16 = 10sec//0.64sec + 1 - self.global_attn_agg = TemporalTransformerEncoderLayer( - add_pos_emb=True, - pos_emb_drop=mformer_cfg.VIT.POS_DROPOUT, - pos_max_len=pos_max_len, - **transf_enc_layer_kwargs) + self.global_attn_agg = TemporalTransformerEncoderLayer(add_pos_emb=True, + pos_emb_drop=mformer_cfg.VIT.POS_DROPOUT, + pos_max_len=pos_max_len, + **transf_enc_layer_kwargs) elif agg_segments_module == 'AveragePooling': self.global_attn_agg = AveragePooling(avg_pattern='B S D -> B D') @@ -222,16 +214,14 @@ class MotionFormer(VisionTransformer): assert self.extract_features # (BS, T, D) where T = 1 + (224 // 16) * (224 // 16) * 8 - x = x[:, - 1:, :] # without the CLS token for efficiency (should be safe for LayerNorm and FC) + x = x[:, 1:, :] # without the CLS token for efficiency (should be safe for LayerNorm and FC) x = self.norm(x) x = self.pre_logits(x) if self.factorize_space_time: x = self.restore_spatio_temp_dims(x, orig_shape) # (B*S, D, t, h, w) <- (B*S, t*h*w, D) x = self.spatial_attn_agg(x, x_mask) # (B*S, t, D) - x = self.temp_attn_agg( - x) # (B*S, D) or (BS, t, D) if `self.temp_attn_agg` is `Identity` + x = self.temp_attn_agg(x) # (B*S, D) or (BS, t, D) if `self.temp_attn_agg` is `Identity` return x @@ -294,8 +284,7 @@ class BaseEncoderLayer(nn.TransformerEncoderLayer): cls_tokens = self.cls_token.expand(batch_dim, -1, -1) # expanding to match batch dimension x = torch.cat((cls_tokens, x), dim=-2) # (batch_dim, 1+seq_len, D) if x_mask is not None: - cls_mask = torch.ones((batch_dim, 1), dtype=torch.bool, - device=x_mask.device) # 1=keep; 0=mask + cls_mask = torch.ones((batch_dim, 1), dtype=torch.bool, device=x_mask.device) # 1=keep; 0=mask x_mask_w_cls = torch.cat((cls_mask, x_mask), dim=-1) # (batch_dim, 1+seq_len) B, N = x_mask_w_cls.shape # torch expects (N, N) or (B*num_heads, N, N) mask (sadness ahead); torch masks @@ -309,8 +298,7 @@ class BaseEncoderLayer(nn.TransformerEncoderLayer): # add positional embedding if self.add_pos_emb: - seq_len = x.shape[ - 1] # (don't even think about moving it before the CLS token concatenation) + seq_len = x.shape[1] # (don't even think about moving it before the CLS token concatenation) assert seq_len <= self.pos_max_len, f'Seq len ({seq_len}) > pos_max_len ({self.pos_max_len})' x = x + self.pos_emb[:, :seq_len, :] x = self.pos_drop(x) diff --git a/fastvideo/third_party/eval/synchformer/resnet.py b/fastvideo/third_party/eval/synchformer/resnet.py index a7f78bd0f..6aad8f9b2 100644 --- a/fastvideo/third_party/eval/synchformer/resnet.py +++ b/fastvideo/third_party/eval/synchformer/resnet.py @@ -39,8 +39,12 @@ class ResNetAudio(ResNet): # replacing the old conv1 to the new one (RGB - 3; spectrogram - 1) conv1 = self.conv1 - self.conv1 = torch.nn.Conv2d(1, conv1.out_channels, conv1.kernel_size, - conv1.stride, conv1.padding, bias=conv1.bias) + self.conv1 = torch.nn.Conv2d(1, + conv1.out_channels, + conv1.kernel_size, + conv1.stride, + conv1.padding, + bias=conv1.bias) self.extract_features = extract_features self.embed_dim = self.fc.in_features @@ -75,20 +79,20 @@ class ResNetAudio(ResNet): class ResNet18AudioFeatures(ResNetAudio): # ckpt_path should default to None, otherwise when no pre-training is desired it will throw an error - def __init__(self, - extract_features: bool = False, - ckpt_path: str = None, - feat_type: str = None, - max_spec_t: int = None, - factorize_freq_time: bool = None, - agg_freq_module: str = None, - agg_time_module: str = None, - add_global_repr: bool = True, - agg_segments_module: str = None, - max_segments: int = None, - ) -> None: - super().__init__(arch_name='resnet18', num_classes=308, extract_features=extract_features, - ckpt_path=ckpt_path) + def __init__( + self, + extract_features: bool = False, + ckpt_path: str = None, + feat_type: str = None, + max_spec_t: int = None, + factorize_freq_time: bool = None, + agg_freq_module: str = None, + agg_time_module: str = None, + add_global_repr: bool = True, + agg_segments_module: str = None, + max_segments: int = None, + ) -> None: + super().__init__(arch_name='resnet18', num_classes=308, extract_features=extract_features, ckpt_path=ckpt_path) assert extract_features, 'Not implemented otherwise' self.extract_features = extract_features self.feat_type = feat_type @@ -113,8 +117,13 @@ class ResNet18AudioFeatures(ResNetAudio): self.factorize_freq_time = factorize_freq_time # avoiding code duplication (used only if agg_*_module is TransformerEncoderLayer) transf_enc_layer_kwargs = dict( - d_model=self.embed_dim, nhead=self.nhead, dim_feedforward=self.mlp_ratio*self.embed_dim, - activation=torch.nn.GELU(), batch_first=True, dropout=self.drop_rate, layer_norm_eps=1e-6, + d_model=self.embed_dim, + nhead=self.nhead, + dim_feedforward=self.mlp_ratio * self.embed_dim, + activation=torch.nn.GELU(), + batch_first=True, + dropout=self.drop_rate, + layer_norm_eps=1e-6, norm_first=True, ) if factorize_freq_time: @@ -139,10 +148,10 @@ class ResNet18AudioFeatures(ResNetAudio): # we can reuse the same layer as for temporal factorization (B, dim_to_agg, D) -> (B, D) # we need to add pos emb (PE) because previously we added the same PE for each segment pos_max_len = max_segments if max_segments is not None else 16 # 16 = 10sec//0.64sec + 1 - self.global_attn_agg = TemporalTransformerEncoderLayer( - add_pos_emb=True, pos_emb_drop=self.drop_rate, - pos_max_len=pos_max_len, **transf_enc_layer_kwargs - ) + self.global_attn_agg = TemporalTransformerEncoderLayer(add_pos_emb=True, + pos_emb_drop=self.drop_rate, + pos_max_len=pos_max_len, + **transf_enc_layer_kwargs) elif agg_segments_module == 'AveragePooling': self.global_attn_agg = AveragePooling(avg_pattern='B S D -> B D') diff --git a/fastvideo/third_party/eval/synchformer/synchformer.py b/fastvideo/third_party/eval/synchformer/synchformer.py index 74e4ead46..9fad69020 100644 --- a/fastvideo/third_party/eval/synchformer/synchformer.py +++ b/fastvideo/third_party/eval/synchformer/synchformer.py @@ -20,13 +20,11 @@ def make_class_grid(leftmost_val, assert grid_size >= 3, f'grid_size: {grid_size} doesnot make sense. If =2 -> (-1,1); =1 -> (-1); =0 -> ()' grid = torch.from_numpy(np.linspace(leftmost_val, rightmost_val, grid_size)).float() if add_extreme_offset: - assert all([seg_size_vframes, nseg, - step_size_seg]), f'{seg_size_vframes} {nseg} {step_size_seg}' + assert all([seg_size_vframes, nseg, step_size_seg]), f'{seg_size_vframes} {nseg} {step_size_seg}' seg_size_sec = seg_size_vframes / vfps trim_size_in_seg = nseg - (1 - step_size_seg) * (nseg - 1) extreme_value = trim_size_in_seg * seg_size_sec - grid = torch.cat([grid, - torch.tensor([extreme_value])]) # adding extreme offset to the class grid + grid = torch.cat([grid, torch.tensor([extreme_value])]) # adding extreme offset to the class grid return grid @@ -68,8 +66,7 @@ class Synchformer(nn.Module): # self.transformer will concatenate the vis and aud in one sequence with aux tokens, # ie `CvvvvMaaaaaa`, and will return the logits for the CLS tokens - logits = self.transformer(vis, - aud) # (B, cls); or (B, cls) and (B, 2) if DoubtingTransformer + logits = self.transformer(vis, aud) # (B, cls); or (B, cls) and (B, 2) if DoubtingTransformer return logits diff --git a/fastvideo/third_party/eval/synchformer/transformer.py b/fastvideo/third_party/eval/synchformer/transformer.py index d12fbd19c..29a9ca1a7 100644 --- a/fastvideo/third_party/eval/synchformer/transformer.py +++ b/fastvideo/third_party/eval/synchformer/transformer.py @@ -9,8 +9,7 @@ from torch.nn import functional as F class GlobalTransformer(torch.nn.Module): '''Same as in SparseSync but without the selector transformers and the head''' - def __init__(self, tok_pdrop, embd_pdrop, resid_pdrop, attn_pdrop, n_layer, n_head, - n_embd) -> None: + def __init__(self, tok_pdrop, embd_pdrop, resid_pdrop, attn_pdrop, n_layer, n_head, n_embd) -> None: super().__init__() self.config = Config(embd_pdrop=embd_pdrop, resid_pdrop=resid_pdrop, @@ -120,18 +119,15 @@ class SelfAttention(nn.Module): # calculate query, key, values for all heads in batch and move head forward to be the batch dim k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) - q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, - 2) # (B, nh, T, hs) - v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, - 2) # (B, nh, T, hs) + q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) # self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T) att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) # att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf')) att = F.softmax(att, dim=-1) y = self.attn_drop(att) @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) - y = y.transpose(1, 2).contiguous().view(B, T, - C) # re-assemble all head outputs side by side + y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side # output projection y = self.resid_drop(self.proj(y)) diff --git a/fastvideo/third_party/eval/synchformer/utils.py b/fastvideo/third_party/eval/synchformer/utils.py index a797eb9c6..60cfc5d42 100644 --- a/fastvideo/third_party/eval/synchformer/utils.py +++ b/fastvideo/third_party/eval/synchformer/utils.py @@ -7,60 +7,37 @@ from tqdm import tqdm PARENT_LINK = 'https://a3s.fi/swift/v1/AUTH_a235c0f452d648828f745589cde1219a' FNAME2LINK = { # S3: Synchability: AudioSet (run 2) - '24-01-22T20-34-52.pt': - f'{PARENT_LINK}/sync/sync_models/24-01-22T20-34-52/24-01-22T20-34-52.pt', - 'cfg-24-01-22T20-34-52.yaml': - f'{PARENT_LINK}/sync/sync_models/24-01-22T20-34-52/cfg-24-01-22T20-34-52.yaml', + '24-01-22T20-34-52.pt': f'{PARENT_LINK}/sync/sync_models/24-01-22T20-34-52/24-01-22T20-34-52.pt', + 'cfg-24-01-22T20-34-52.yaml': f'{PARENT_LINK}/sync/sync_models/24-01-22T20-34-52/cfg-24-01-22T20-34-52.yaml', # S2: Synchformer: AudioSet (run 2) - '24-01-04T16-39-21.pt': - f'{PARENT_LINK}/sync/sync_models/24-01-04T16-39-21/24-01-04T16-39-21.pt', - 'cfg-24-01-04T16-39-21.yaml': - f'{PARENT_LINK}/sync/sync_models/24-01-04T16-39-21/cfg-24-01-04T16-39-21.yaml', + '24-01-04T16-39-21.pt': f'{PARENT_LINK}/sync/sync_models/24-01-04T16-39-21/24-01-04T16-39-21.pt', + 'cfg-24-01-04T16-39-21.yaml': f'{PARENT_LINK}/sync/sync_models/24-01-04T16-39-21/cfg-24-01-04T16-39-21.yaml', # S2: Synchformer: AudioSet (run 1) - '23-08-28T11-23-23.pt': - f'{PARENT_LINK}/sync/sync_models/23-08-28T11-23-23/23-08-28T11-23-23.pt', - 'cfg-23-08-28T11-23-23.yaml': - f'{PARENT_LINK}/sync/sync_models/23-08-28T11-23-23/cfg-23-08-28T11-23-23.yaml', + '23-08-28T11-23-23.pt': f'{PARENT_LINK}/sync/sync_models/23-08-28T11-23-23/23-08-28T11-23-23.pt', + 'cfg-23-08-28T11-23-23.yaml': f'{PARENT_LINK}/sync/sync_models/23-08-28T11-23-23/cfg-23-08-28T11-23-23.yaml', # S2: Synchformer: LRS3 (run 2) - '23-12-23T18-33-57.pt': - f'{PARENT_LINK}/sync/sync_models/23-12-23T18-33-57/23-12-23T18-33-57.pt', - 'cfg-23-12-23T18-33-57.yaml': - f'{PARENT_LINK}/sync/sync_models/23-12-23T18-33-57/cfg-23-12-23T18-33-57.yaml', + '23-12-23T18-33-57.pt': f'{PARENT_LINK}/sync/sync_models/23-12-23T18-33-57/23-12-23T18-33-57.pt', + 'cfg-23-12-23T18-33-57.yaml': f'{PARENT_LINK}/sync/sync_models/23-12-23T18-33-57/cfg-23-12-23T18-33-57.yaml', # S2: Synchformer: VGS (run 2) - '24-01-02T10-00-53.pt': - f'{PARENT_LINK}/sync/sync_models/24-01-02T10-00-53/24-01-02T10-00-53.pt', - 'cfg-24-01-02T10-00-53.yaml': - f'{PARENT_LINK}/sync/sync_models/24-01-02T10-00-53/cfg-24-01-02T10-00-53.yaml', + '24-01-02T10-00-53.pt': f'{PARENT_LINK}/sync/sync_models/24-01-02T10-00-53/24-01-02T10-00-53.pt', + 'cfg-24-01-02T10-00-53.yaml': f'{PARENT_LINK}/sync/sync_models/24-01-02T10-00-53/cfg-24-01-02T10-00-53.yaml', # SparseSync: ft VGGSound-Full - '22-09-21T21-00-52.pt': - f'{PARENT_LINK}/sync/sync_models/22-09-21T21-00-52/22-09-21T21-00-52.pt', - 'cfg-22-09-21T21-00-52.yaml': - f'{PARENT_LINK}/sync/sync_models/22-09-21T21-00-52/cfg-22-09-21T21-00-52.yaml', + '22-09-21T21-00-52.pt': f'{PARENT_LINK}/sync/sync_models/22-09-21T21-00-52/22-09-21T21-00-52.pt', + 'cfg-22-09-21T21-00-52.yaml': f'{PARENT_LINK}/sync/sync_models/22-09-21T21-00-52/cfg-22-09-21T21-00-52.yaml', # SparseSync: ft VGGSound-Sparse - '22-07-28T15-49-45.pt': - f'{PARENT_LINK}/sync/sync_models/22-07-28T15-49-45/22-07-28T15-49-45.pt', - 'cfg-22-07-28T15-49-45.yaml': - f'{PARENT_LINK}/sync/sync_models/22-07-28T15-49-45/cfg-22-07-28T15-49-45.yaml', + '22-07-28T15-49-45.pt': f'{PARENT_LINK}/sync/sync_models/22-07-28T15-49-45/22-07-28T15-49-45.pt', + 'cfg-22-07-28T15-49-45.yaml': f'{PARENT_LINK}/sync/sync_models/22-07-28T15-49-45/cfg-22-07-28T15-49-45.yaml', # SparseSync: only pt on LRS3 - '22-07-13T22-25-49.pt': - f'{PARENT_LINK}/sync/sync_models/22-07-13T22-25-49/22-07-13T22-25-49.pt', - 'cfg-22-07-13T22-25-49.yaml': - f'{PARENT_LINK}/sync/sync_models/22-07-13T22-25-49/cfg-22-07-13T22-25-49.yaml', + '22-07-13T22-25-49.pt': f'{PARENT_LINK}/sync/sync_models/22-07-13T22-25-49/22-07-13T22-25-49.pt', + 'cfg-22-07-13T22-25-49.yaml': f'{PARENT_LINK}/sync/sync_models/22-07-13T22-25-49/cfg-22-07-13T22-25-49.yaml', # SparseSync: feature extractors - 'ResNetAudio-22-08-04T09-51-04.pt': - f'{PARENT_LINK}/sync/ResNetAudio-22-08-04T09-51-04.pt', # 2s - 'ResNetAudio-22-08-03T23-14-49.pt': - f'{PARENT_LINK}/sync/ResNetAudio-22-08-03T23-14-49.pt', # 3s - 'ResNetAudio-22-08-03T23-14-28.pt': - f'{PARENT_LINK}/sync/ResNetAudio-22-08-03T23-14-28.pt', # 4s - 'ResNetAudio-22-06-24T08-10-33.pt': - f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T08-10-33.pt', # 5s - 'ResNetAudio-22-06-24T17-31-07.pt': - f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T17-31-07.pt', # 6s - 'ResNetAudio-22-06-24T23-57-11.pt': - f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T23-57-11.pt', # 7s - 'ResNetAudio-22-06-25T04-35-42.pt': - f'{PARENT_LINK}/sync/ResNetAudio-22-06-25T04-35-42.pt', # 8s + 'ResNetAudio-22-08-04T09-51-04.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-08-04T09-51-04.pt', # 2s + 'ResNetAudio-22-08-03T23-14-49.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-08-03T23-14-49.pt', # 3s + 'ResNetAudio-22-08-03T23-14-28.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-08-03T23-14-28.pt', # 4s + 'ResNetAudio-22-06-24T08-10-33.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T08-10-33.pt', # 5s + 'ResNetAudio-22-06-24T17-31-07.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T17-31-07.pt', # 6s + 'ResNetAudio-22-06-24T23-57-11.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T23-57-11.pt', # 7s + 'ResNetAudio-22-06-25T04-35-42.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-06-25T04-35-42.pt', # 8s } diff --git a/fastvideo/third_party/eval/synchformer/video_model_builder.py b/fastvideo/third_party/eval/synchformer/video_model_builder.py index 75e6d712c..f0e47daeb 100644 --- a/fastvideo/third_party/eval/synchformer/video_model_builder.py +++ b/fastvideo/third_party/eval/synchformer/video_model_builder.py @@ -55,8 +55,7 @@ class VisionTransformer(nn.Module): in_chans=self.in_chans, embed_dim=self.embed_dim, z_block_size=self.cfg.VIT.PATCH_SIZE_TEMP) - self.patch_embed_3d.proj.weight.data = torch.zeros_like( - self.patch_embed_3d.proj.weight.data) + self.patch_embed_3d.proj.weight.data = torch.zeros_like(self.patch_embed_3d.proj.weight.data) # Number of patches if self.video_input: @@ -70,8 +69,7 @@ class VisionTransformer(nn.Module): trunc_normal_(self.cls_token, std=.02) # Positional embedding - self.pos_embed = nn.Parameter( - torch.zeros(1, self.patch_embed.num_patches + 1, self.embed_dim)) + self.pos_embed = nn.Parameter(torch.zeros(1, self.patch_embed.num_patches + 1, self.embed_dim)) self.pos_drop = nn.Dropout(p=cfg.VIT.POS_DROPOUT) trunc_normal_(self.pos_embed, std=.02) @@ -108,8 +106,7 @@ class VisionTransformer(nn.Module): attn_drop=self.attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, - use_original_code=self.cfg.VIT.USE_ORIGINAL_TRAJ_ATTN_CODE) - for i in range(self.depth) + use_original_code=self.cfg.VIT.USE_ORIGINAL_TRAJ_ATTN_CODE) for i in range(self.depth) ]) self.norm = norm_layer(self.embed_dim) @@ -125,11 +122,10 @@ class VisionTransformer(nn.Module): else: # logging.info("Using ReLU activation in MLP") act = nn.ReLU() - self.pre_logits = nn.Sequential( - OrderedDict([ - ('fc', nn.Linear(self.embed_dim, hidden_dim)), - ('act', act), - ])) + self.pre_logits = nn.Sequential(OrderedDict([ + ('fc', nn.Linear(self.embed_dim, hidden_dim)), + ('act', act), + ])) else: self.pre_logits = nn.Identity() @@ -139,8 +135,7 @@ class VisionTransformer(nn.Module): for a, i in enumerate(range(len(self.num_classes))): setattr(self, "head%d" % a, nn.Linear(self.embed_dim, self.num_classes[i])) else: - self.head = nn.Linear(self.embed_dim, - self.num_classes) if self.num_classes > 0 else nn.Identity() + self.head = nn.Linear(self.embed_dim, self.num_classes) if self.num_classes > 0 else nn.Identity() # Initialize weights self.apply(self._init_weights) diff --git a/fastvideo/third_party/eval/synchformer/vit_helper.py b/fastvideo/third_party/eval/synchformer/vit_helper.py index 6af730a13..a4165222d 100644 --- a/fastvideo/third_party/eval/synchformer/vit_helper.py +++ b/fastvideo/third_party/eval/synchformer/vit_helper.py @@ -26,8 +26,7 @@ def qkv_attn(q, k, v, tok_mask: torch.Tensor = None): # apply masking if provided, tok_mask is (B*S*H, N): 1s - keep; sim is (B*S*H, H, N, N) if tok_mask is not None: BSH, N = tok_mask.shape - sim = sim.masked_fill(tok_mask.view(BSH, 1, N) == 0, - float('-inf')) # 1 - broadcasts across N + sim = sim.masked_fill(tok_mask.view(BSH, 1, N) == 0, float('-inf')) # 1 - broadcasts across N attn = sim.softmax(dim=-1) out = einsum('b i j, b j d -> b i d', attn, v) return out @@ -79,8 +78,7 @@ class DividedAttention(nn.Module): cls_out = qkv_attn(cls_q, k, v, tok_mask=tok_mask) # rearrange across time or space - q_, k_, v_ = map(lambda t: rearrange(t, f'{einops_from} -> {einops_to}', **einops_dims), - (q_, k_, v_)) + q_, k_, v_ = map(lambda t: rearrange(t, f'{einops_from} -> {einops_to}', **einops_dims), (q_, k_, v_)) # expand CLS token keys and values across time or space and concat r = q_.shape[0] // cls_k.shape[0] @@ -92,10 +90,8 @@ class DividedAttention(nn.Module): # the same for masking (if provided) if tok_mask is not None: # since mask does not have the latent dim (d), we need to remove it from einops dims - mask_ = rearrange(mask_, f'{einops_from} -> {einops_to}'.replace(' d', ''), - **einops_dims) - cls_mask = repeat(cls_mask, 'b () -> (b r) ()', - r=r) # expand cls_mask across time or space + mask_ = rearrange(mask_, f'{einops_from} -> {einops_to}'.replace(' d', ''), **einops_dims) + cls_mask = repeat(cls_mask, 'b () -> (b r) ()', r=r) # expand cls_mask across time or space mask_ = torch.cat((cls_mask, mask_), dim=1) # attention @@ -138,11 +134,7 @@ class DividedSpaceTimeBlock(nn.Module): self.norm1 = norm_layer(dim) - self.attn = DividedAttention(dim, - num_heads=num_heads, - qkv_bias=qkv_bias, - attn_drop=attn_drop, - proj_drop=drop) + self.attn = DividedAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) self.timeattn = DividedAttention(dim, num_heads=num_heads, @@ -154,19 +146,10 @@ class DividedSpaceTimeBlock(nn.Module): self.drop_path = nn.Identity() self.norm2 = norm_layer(dim) mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp(in_features=dim, - hidden_features=mlp_hidden_dim, - act_layer=act_layer, - drop=drop) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) self.norm3 = norm_layer(dim) - def forward(self, - x, - seq_len=196, - num_frames=8, - approx='none', - num_landmarks=128, - tok_mask: torch.Tensor = None): + def forward(self, x, seq_len=196, num_frames=8, approx='none', num_landmarks=128, tok_mask: torch.Tensor = None): time_output = self.timeattn(self.norm3(x), self.einops_from_time, self.einops_to_time, @@ -188,12 +171,7 @@ class DividedSpaceTimeBlock(nn.Module): class Mlp(nn.Module): - def __init__(self, - in_features, - hidden_features=None, - out_features=None, - act_layer=nn.GELU, - drop=0.): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features @@ -274,14 +252,11 @@ class HeadMLP(nn.Module): self.n_hidden = n_hidden if n_hidden is None: # use linear classifier - self.block_forward = nn.Sequential(nn.Dropout(p=p), - nn.Linear(n_input, n_classes, bias=True)) + self.block_forward = nn.Sequential(nn.Dropout(p=p), nn.Linear(n_input, n_classes, bias=True)) else: # use simple MLP classifier - self.block_forward = nn.Sequential(nn.Dropout(p=p), - nn.Linear(n_input, n_hidden, bias=True), - nn.BatchNorm1d(n_hidden), nn.ReLU(inplace=True), - nn.Dropout(p=p), + self.block_forward = nn.Sequential(nn.Dropout(p=p), nn.Linear(n_input, n_hidden, bias=True), + nn.BatchNorm1d(n_hidden), nn.ReLU(inplace=True), nn.Dropout(p=p), nn.Linear(n_hidden, n_classes, bias=True)) print(f"Dropout-NLP: {p}") @@ -333,13 +308,7 @@ def adapt_input_conv(in_chans, conv_weight, agg='sum'): return conv_weight -def load_pretrained(model, - cfg=None, - num_classes=1000, - in_chans=3, - filter_fn=None, - strict=True, - progress=False): +def load_pretrained(model, cfg=None, num_classes=1000, in_chans=3, filter_fn=None, strict=True, progress=False): # Load state dict assert (f"{cfg.VIT.PRETRAINED_WEIGHTS} not in [vit_1k, vit_1k_large]") state_dict = torch.hub.load_state_dict_from_url(url=default_cfgs[cfg.VIT.PRETRAINED_WEIGHTS]) @@ -354,18 +323,12 @@ def load_pretrained(model, for input_conv_name in input_convs: weight_name = input_conv_name + '.weight' try: - state_dict[weight_name] = adapt_input_conv(in_chans, - state_dict[weight_name], - agg='avg') - print( - f'Converted input conv {input_conv_name} pretrained weights from 3 to {in_chans} channel(s)' - ) + state_dict[weight_name] = adapt_input_conv(in_chans, state_dict[weight_name], agg='avg') + print(f'Converted input conv {input_conv_name} pretrained weights from 3 to {in_chans} channel(s)') except NotImplementedError as e: del state_dict[weight_name] strict = False - print( - f'Unable to convert pretrained {input_conv_name} weights, using random init for this layer.' - ) + print(f'Unable to convert pretrained {input_conv_name} weights, using random init for this layer.') classifier_name = 'head' label_offset = cfg.get('label_offset', 0) diff --git a/fastvideo/third_party/longcat_video/block_sparse_attention/bsa_interface.py b/fastvideo/third_party/longcat_video/block_sparse_attention/bsa_interface.py index 8e9e13f5a..1b6bbd6c9 100644 --- a/fastvideo/third_party/longcat_video/block_sparse_attention/bsa_interface.py +++ b/fastvideo/third_party/longcat_video/block_sparse_attention/bsa_interface.py @@ -6,19 +6,23 @@ import triton.language as tl import math from .common import _attn_fwd_gating, _attn_bwd_preprocess, configs_gating_preset -from .flash_attn_bsa_varlen_mask import ( - _attn_fwd_bsa_varlen, _attn_fwd_bsa_varlen_align, _attn_bwd_dkdv_bsa_varlen_wrapper, _attn_bwd_dq_bsa_varlen_wrapper, _attn_bwd_dq_bsa_varlen_align_wrapper, - configs_fwd_bsa_varlen_preset, configs_fwd_bsa_varlen_align_preset, configs_bwd_dkdv_bsa_varlen_preset, configs_bwd_dq_bsa_varlen_preset, configs_bwd_dq_bsa_varlen_align_preset -) +from .flash_attn_bsa_varlen_mask import (_attn_fwd_bsa_varlen, _attn_fwd_bsa_varlen_align, + _attn_bwd_dkdv_bsa_varlen_wrapper, _attn_bwd_dq_bsa_varlen_wrapper, + _attn_bwd_dq_bsa_varlen_align_wrapper, configs_fwd_bsa_varlen_preset, + configs_fwd_bsa_varlen_align_preset, configs_bwd_dkdv_bsa_varlen_preset, + configs_bwd_dq_bsa_varlen_preset, configs_bwd_dq_bsa_varlen_align_preset) torch._dynamo.config.cache_size_limit = 32 + def is_cuda(): return triton.runtime.driver.active.get_current_target().backend == "cuda" + def supports_tma(): return is_cuda() and torch.cuda.get_device_capability()[0] >= 9 + HAS_TMA_DESC = "nv_tma_desc_type" in dir(tl) if HAS_TMA_DESC: @@ -93,61 +97,75 @@ class TmaAutoTuneHelper: def create_mask_from_indices_kernel( block_indices, block_mask, - stride_bz, stride_bh, stride_bm, stride_bs, - stride_mz, stride_mh, stride_mm, stride_mn, + stride_bz, + stride_bh, + stride_bm, + stride_bs, + stride_mz, + stride_mh, + stride_mm, + stride_mn, H, ): i_zh, i_m, i_s = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_z, i_h = i_zh // H, i_zh % H - off_b = i_z.to(tl.int64) * stride_bz + i_h.to(tl.int64) * stride_bh + i_m.to(tl.int64) * stride_bm + i_s.to(tl.int64) * stride_bs - + off_b = i_z.to(tl.int64) * stride_bz + i_h.to(tl.int64) * stride_bh + i_m.to(tl.int64) * stride_bm + i_s.to( + tl.int64) * stride_bs + b_i = tl.load(block_indices + off_b) - - off_m = i_z.to(tl.int64) * stride_mz + i_h.to(tl.int64) * stride_mh + i_m.to(tl.int64) * stride_mm + b_i.to(tl.int64) * stride_mn - + + off_m = i_z.to(tl.int64) * stride_mz + i_h.to(tl.int64) * stride_mh + i_m.to(tl.int64) * stride_mm + b_i.to( + tl.int64) * stride_mn + b_m = 1 tl.store(block_mask + off_m, b_m.to(block_mask.dtype.element_ty)) -def create_mask_from_indices_triton( - block_indices, - N_cols -): + +def create_mask_from_indices_triton(block_indices, N_cols): B, H, N_rows, S = block_indices.shape block_mask = torch.zeros((B, H, N_rows, N_cols), dtype=torch.bool, device=block_indices.device) create_mask_from_indices_kernel[(B * H, N_rows, S)]( block_indices, block_mask, - block_indices.stride(0), block_indices.stride(1), block_indices.stride(2), block_indices.stride(3), - block_mask.stride(0), block_mask.stride(1), block_mask.stride(2), block_mask.stride(3), + block_indices.stride(0), + block_indices.stride(1), + block_indices.stride(2), + block_indices.stride(3), + block_mask.stride(0), + block_mask.stride(1), + block_mask.stride(2), + block_mask.stride(3), H, ) return block_mask + @torch.compile def create_mask_from_indices_varlen(block_indices, N_cols_mask): - + B, H, M, _ = block_indices.shape device = block_indices.device - + mask = torch.zeros((B, H, M, N_cols_mask), dtype=torch.bool, device=device) - + valid = block_indices < N_cols_mask - + b_idx = torch.arange(B, device=device)[:, None, None, None].expand_as(block_indices) h_idx = torch.arange(H, device=device)[None, :, None, None].expand_as(block_indices) m_idx = torch.arange(M, device=device)[None, None, :, None].expand_as(block_indices) - + valid_coords = (b_idx[valid], h_idx[valid], m_idx[valid], block_indices[valid]) - + mask[valid_coords] = True - + return mask + @torch.compile def create_indices_k_from_indices_q_varlen( block_indices, - N_cols_mask # indicate the number of the last dimension of the bool mask, since this information cannot be determined by block_indices, which may contain invalid elements + N_cols_mask # indicate the number of the last dimension of the bool mask, since this information cannot be determined by block_indices, which may contain invalid elements ): block_mask_qk = create_mask_from_indices_varlen(block_indices, N_cols_mask) B, H, M, N = block_mask_qk.shape @@ -155,17 +173,14 @@ def create_indices_k_from_indices_q_varlen( indices = torch.arange(M, device=block_indices.device).view(1, 1, 1, -1).expand_as(block_mask_kq) block_indices_k = torch.where(block_mask_kq, indices, M) block_indices_k, _ = torch.sort(block_indices_k, dim=-1) - + block_indices_k_lens = (block_indices_k < M).sum(dim=-1) return block_indices_k, block_indices_k_lens @torch.compile -def mean_pooling_compression( - x: torch.Tensor, - block_size: int -) -> torch.Tensor: +def mean_pooling_compression(x: torch.Tensor, block_size: int) -> torch.Tensor: B, H, S = x.shape[:3] num_block = math.ceil(S / block_size) if S % block_size != 0: @@ -173,52 +188,66 @@ def mean_pooling_compression( x_cmp = x.view(B, H, num_block, block_size, -1).mean(dim=3) return x_cmp + @torch.compile def cal_score(q, k): k_transposed = k.transpose(-1, -2) # [b, h, d, s_k] score = torch.matmul(q, k_transposed) # [b, h, s_q, s_k] return score + def cal_score_triton(q, k): B, H, s_q, D = q.shape s_k = k.shape[2] - + score = torch.empty(B, H, s_q, s_k, device=q.device, dtype=q.dtype) - + kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_gating_preset['default'] - + grid = lambda args: (triton.cdiv(s_q, args["BLOCK_M"]), B * H, 1) - _attn_fwd_gating[grid]( - q, k, score, - q.stride(0), q.stride(1), q.stride(2), q.stride(3), - k.stride(0), k.stride(1), k.stride(2), k.stride(3), - score.stride(0), score.stride(1), score.stride(2), score.stride(3), - H, s_q, s_k, - HEAD_DIM=D, - **kernel_config - ) + _attn_fwd_gating[grid](q, + k, + score, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + score.stride(0), + score.stride(1), + score.stride(2), + score.stride(3), + H, + s_q, + s_k, + HEAD_DIM=D, + **kernel_config) return score + @torch.compile def get_select_indices_topk(q, k, sparsity): score = cal_score(q, k) block_indices, block_indices_lens = get_select_indices_topk_from_score(score, sparsity) return block_indices, block_indices_lens + @torch.compile def get_select_indices_topk_from_score(score, sparsity): num_selected = int((1 - sparsity) * score.shape[-1]) block_indices = torch.topk(score, num_selected)[1] - block_indices_lens = torch.full( - (block_indices.shape[0], block_indices.shape[1], block_indices.shape[2]), - num_selected, - dtype=torch.int32, - device=block_indices.device - ) + block_indices_lens = torch.full((block_indices.shape[0], block_indices.shape[1], block_indices.shape[2]), + num_selected, + dtype=torch.int32, + device=block_indices.device) return block_indices, block_indices_lens + @torch.compile def get_select_indices_cdf(q, k, cdf_threshold): score = cal_score(q, k) @@ -226,31 +255,35 @@ def get_select_indices_cdf(q, k, cdf_threshold): block_indices, block_indices_lens = get_select_indices_cdf_from_score(score, cdf_threshold, 1 / head_dim**0.5) return block_indices, block_indices_lens + @torch.compile def get_select_indices_cdf_from_score(score, cdf_threshold, sm_scale): weights = torch.softmax(score * sm_scale, dim=-1) - + B, H, Sq, Sk = weights.shape - cdf_threshold = torch.full((H,), cdf_threshold, device=weights.device).view(1, H, 1, 1).expand(B, -1, Sq, -1) + cdf_threshold = torch.full((H, ), cdf_threshold, device=weights.device).view(1, H, 1, 1).expand(B, -1, Sq, -1) weights_sorted = torch.sort(weights, dim=-1, descending=True) cdf = torch.cumsum(weights_sorted.values, dim=-1) num_selected = torch.searchsorted(cdf, cdf_threshold, right=True) - + return weights_sorted.indices, num_selected.squeeze(-1) + @torch.compile def get_select_indices_cdf_topk(q, k, sparsity, cdf_threshold): score = cal_score(q, k) head_dim = q.shape[-1] - block_indices, block_indices_lens = get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold, 1 / head_dim**0.5) + block_indices, block_indices_lens = get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold, + 1 / head_dim**0.5) return block_indices, block_indices_lens + @torch.compile def get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold, sm_scale): weights = torch.softmax(score * sm_scale, dim=-1) - + B, H, Sq, Sk = weights.shape - cdf_threshold = torch.full((H,), cdf_threshold, device=weights.device).view(1, H, 1, 1).expand(B, -1, Sq, -1) + cdf_threshold = torch.full((H, ), cdf_threshold, device=weights.device).view(1, H, 1, 1).expand(B, -1, Sq, -1) weights_sorted = torch.sort(weights, dim=-1, descending=True) cdf = torch.cumsum(weights_sorted.values, dim=-1) num_selected = torch.searchsorted(cdf, cdf_threshold, right=True) @@ -261,6 +294,7 @@ def get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold, sm_sc return weights_sorted.indices, num_selected.squeeze(-1) + def get_select_indices(q, k, sparsity, cdf_threshold): if sparsity is not None and cdf_threshold is None: block_indices, block_indices_lens = get_select_indices_topk(q, k, sparsity) @@ -272,6 +306,7 @@ def get_select_indices(q, k, sparsity, cdf_threshold): raise ValueError return block_indices, block_indices_lens + def get_select_indices_from_score(score, sparsity, cdf_threshold): if sparsity is not None and cdf_threshold is None: block_indices, block_indices_lens = get_select_indices_topk_from_score(score, sparsity) @@ -283,175 +318,230 @@ def get_select_indices_from_score(score, sparsity, cdf_threshold): raise ValueError return block_indices, block_indices_lens -def attn_fwd_bsa_varlen_triton( - q, - k, - v, - sm_scale, - block_indices, - block_indices_lens, - chunk_size_q, - chunk_size_k, - sparsity -): - + +def attn_fwd_bsa_varlen_triton(q, k, v, sm_scale, block_indices, block_indices_lens, chunk_size_q, chunk_size_k, + sparsity): + B, H, Seq, D = q.shape o = torch.empty_like(q) M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) - + grid = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), q.shape[0] * q.shape[1], 1) config_key = 'BLOCK_N_LG=64' if chunk_size_k == 64 else 'default' if chunk_size_k > 128: fwd_func = _attn_fwd_bsa_varlen - kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_fwd_bsa_varlen_preset[config_key] + kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', + '0') == '1' else configs_fwd_bsa_varlen_preset[config_key] else: fwd_func = _attn_fwd_bsa_varlen_align - kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_fwd_bsa_varlen_align_preset[config_key] - + kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', + '0') == '1' else configs_fwd_bsa_varlen_align_preset[config_key] + block_indices = block_indices.contiguous() block_indices_lens = block_indices_lens.contiguous() - + fwd_func[grid]( - q, k, v, sm_scale, M, o, - block_indices, # [B, H, M_COMPRESS, S] - block_indices_lens, # [B, H, M_COMPRESS, S_MAX] - q.stride(0), q.stride(1), q.stride(2), q.stride(3), - k.stride(0), k.stride(1), k.stride(2), k.stride(3), - v.stride(0), v.stride(1), v.stride(2), v.stride(3), - o.stride(0), o.stride(1), o.stride(2), o.stride(3), - block_indices.stride(0), block_indices.stride(1), block_indices.stride(2), block_indices.stride(3), - block_indices_lens.stride(0), block_indices_lens.stride(1), block_indices_lens.stride(2), - H, Seq, + q, + k, + v, + sm_scale, + M, + o, + block_indices, # [B, H, M_COMPRESS, S] + block_indices_lens, # [B, H, M_COMPRESS, S_MAX] + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + o.stride(0), + o.stride(1), + o.stride(2), + o.stride(3), + block_indices.stride(0), + block_indices.stride(1), + block_indices.stride(2), + block_indices.stride(3), + block_indices_lens.stride(0), + block_indices_lens.stride(1), + block_indices_lens.stride(2), + H, + Seq, D, BLOCK_M=chunk_size_q, BLOCK_N_LG=chunk_size_k, SPARSITY=sparsity, - **kernel_config - ) - + **kernel_config) + LN2 = 0.6931471824645996 - lse = M * LN2 # convert back to natural units (M is of base 2) + lse = M * LN2 # convert back to natural units (M is of base 2) return o, lse -def attn_bwd_bsa_varlen_triton( - do, - q, - k, - v, - o, - dq, - dk, - dv, - sm_scale, - M, - block_indices, - block_indices_lens, - chunk_size_q, - chunk_size_k, - sparsity -): + +def attn_bwd_bsa_varlen_triton(do, q, k, v, o, dq, dk, dv, sm_scale, M, block_indices, block_indices_lens, chunk_size_q, + chunk_size_k, sparsity): RCP_LN2 = 1.4426950408889634 - M = M * RCP_LN2 # ln -> log2 - + M = M * RCP_LN2 # ln -> log2 + do = do.contiguous() # assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride() BATCH, N_HEAD, N_CTX, HEAD_DIM = q.shape N_CTX_KV = k.shape[-2] - RCP_LN2 = 1.4426950408889634 # = 1.0 / ln(2) # reciprocal + RCP_LN2 = 1.4426950408889634 # = 1.0 / ln(2) # reciprocal arg_k = k arg_k = arg_k * (sm_scale * RCP_LN2) - + PRE_BLOCK = 128 if min(chunk_size_q, chunk_size_k) >= 128 else min(chunk_size_q, chunk_size_k) - + assert N_CTX % PRE_BLOCK == 0 pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD) delta = torch.empty_like(M) - _attn_bwd_preprocess[pre_grid]( - o, do, - delta, - N_CTX, - BLOCK_M=PRE_BLOCK, - HEAD_DIM=HEAD_DIM - ) + _attn_bwd_preprocess[pre_grid](o, do, delta, N_CTX, BLOCK_M=PRE_BLOCK, HEAD_DIM=HEAD_DIM) + + block_indices_k, block_indices_k_lens = create_indices_k_from_indices_q_varlen(block_indices=block_indices, + N_cols_mask=N_CTX_KV // chunk_size_k) - block_indices_k, block_indices_k_lens = create_indices_k_from_indices_q_varlen( - block_indices=block_indices, - N_cols_mask=N_CTX_KV // chunk_size_k - ) - block_indices = block_indices.contiguous() block_indices_lens = block_indices_lens.contiguous() block_indices_k = block_indices_k.contiguous() block_indices_k_lens = block_indices_k_lens.contiguous() config_key = 'BLOCK_N_DQ_LG=64' if chunk_size_k == 64 else 'default' - kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_bwd_dkdv_bsa_varlen_preset[config_key] - + kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', + '0') == '1' else configs_bwd_dkdv_bsa_varlen_preset[config_key] + grid_dkdv = lambda args: (triton.cdiv(arg_k.shape[2], args["BLOCK_N"]), 1, arg_k.shape[0] * arg_k.shape[1]) _attn_bwd_dkdv_bsa_varlen_wrapper[grid_dkdv]( - q, arg_k, v, sm_scale, # softmax scale + q, + arg_k, + v, + sm_scale, # softmax scale do, - dk, dv, - M, # lse (log2) + dk, + dv, + M, # lse (log2) delta, block_indices_k, block_indices_k_lens, - q.stride(0), q.stride(1), q.stride(2), q.stride(3), - k.stride(0), k.stride(1), k.stride(2), k.stride(3), - v.stride(0), v.stride(1), v.stride(2), v.stride(3), - dk.stride(0), dk.stride(1), dk.stride(2), dk.stride(3), - dv.stride(0), dv.stride(1), dv.stride(2), dv.stride(3), - do.stride(0), do.stride(1), do.stride(2), do.stride(3), - M.stride(0), M.stride(1), M.stride(2), - delta.stride(0), delta.stride(1), delta.stride(2), - block_indices_k.stride(0), block_indices_k.stride(1), block_indices_k.stride(2), block_indices_k.stride(3), - block_indices_k_lens.stride(0), block_indices_k_lens.stride(1), block_indices_k_lens.stride(2), - N_HEAD, N_CTX, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + dk.stride(0), + dk.stride(1), + dk.stride(2), + dk.stride(3), + dv.stride(0), + dv.stride(1), + dv.stride(2), + dv.stride(3), + do.stride(0), + do.stride(1), + do.stride(2), + do.stride(3), + M.stride(0), + M.stride(1), + M.stride(2), + delta.stride(0), + delta.stride(1), + delta.stride(2), + block_indices_k.stride(0), + block_indices_k.stride(1), + block_indices_k.stride(2), + block_indices_k.stride(3), + block_indices_k_lens.stride(0), + block_indices_k_lens.stride(1), + block_indices_k_lens.stride(2), + N_HEAD, + N_CTX, BLOCK_M=chunk_size_q, BLOCK_N_DQ_LG=chunk_size_k, HEAD_DIM=HEAD_DIM, SPARSITY=sparsity, - **kernel_config - ) + **kernel_config) config_key = 'BLOCK_N_DQ_LG=64' if chunk_size_k == 64 else 'default' if chunk_size_k > 128: bwd_dq_func = _attn_bwd_dq_bsa_varlen_wrapper - kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_bwd_dq_bsa_varlen_preset[config_key] + kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', + '0') == '1' else configs_bwd_dq_bsa_varlen_preset[config_key] else: bwd_dq_func = _attn_bwd_dq_bsa_varlen_align_wrapper - kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_bwd_dq_bsa_varlen_align_preset[config_key] - + kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', + '0') == '1' else configs_bwd_dq_bsa_varlen_align_preset[config_key] + grid_dq = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), 1, q.shape[0] * q.shape[1]) bwd_dq_func[grid_dq]( - q, arg_k, v, - do, + q, + arg_k, + v, + do, dq, - M, # lse (log2) + M, # lse (log2) delta, block_indices, block_indices_lens, - q.stride(0), q.stride(1), q.stride(2), q.stride(3), - k.stride(0), k.stride(1), k.stride(2), k.stride(3), - v.stride(0), v.stride(1), v.stride(2), v.stride(3), - dq.stride(0), dq.stride(1), dq.stride(2), dq.stride(3), - do.stride(0), do.stride(1), do.stride(2), do.stride(3), - M.stride(0), M.stride(1), M.stride(2), - delta.stride(0), delta.stride(1), delta.stride(2), - block_indices.stride(0), block_indices.stride(1), block_indices.stride(2), block_indices.stride(3), - block_indices_lens.stride(0), block_indices_lens.stride(1), block_indices_lens.stride(2), - N_HEAD, N_CTX, - BLOCK_M=chunk_size_q, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + dq.stride(0), + dq.stride(1), + dq.stride(2), + dq.stride(3), + do.stride(0), + do.stride(1), + do.stride(2), + do.stride(3), + M.stride(0), + M.stride(1), + M.stride(2), + delta.stride(0), + delta.stride(1), + delta.stride(2), + block_indices.stride(0), + block_indices.stride(1), + block_indices.stride(2), + block_indices.stride(3), + block_indices_lens.stride(0), + block_indices_lens.stride(1), + block_indices_lens.stride(2), + N_HEAD, + N_CTX, + BLOCK_M=chunk_size_q, BLOCK_N_DQ_LG=chunk_size_k, HEAD_DIM=HEAD_DIM, SPARSITY=sparsity, - **kernel_config - ) + **kernel_config) + @torch.compile def make_block_indices_varlen_cp_list(block_indices, cp_size, num_blocks_k_full): @@ -470,33 +560,40 @@ def make_block_indices_varlen_cp_list(block_indices, cp_size, num_blocks_k_full) block_indices_tmp = block_indices.clone() min_block_idx = i * num_blocks_per_rank block_indices_tmp -= min_block_idx - block_indices_tmp[block_indices_tmp < 0] = num_blocks_per_rank # block_indices_tmp < 0 indicate invalid indices, set them to num_blocks_per_rank in order to sort them to the tail, so that the first N elements of the block_indices indicated by block_indices_lens are valid - + block_indices_tmp[ + block_indices_tmp < + 0] = num_blocks_per_rank # block_indices_tmp < 0 indicate invalid indices, set them to num_blocks_per_rank in order to sort them to the tail, so that the first N elements of the block_indices indicated by block_indices_lens are valid + block_indices_tmp, _ = torch.sort(block_indices_tmp, dim=-1) - + block_indices_tmp_lens = (block_indices_tmp < num_blocks_per_rank).sum(dim=-1) - + res.append([block_indices_tmp, block_indices_tmp_lens]) return res + @torch.compile def flash_attn_fwd_softmax_lse_correction( - softmax_lse: torch.Tensor, + softmax_lse: torch.Tensor, softmax_lse_per_step: torch.Tensor, ): """Merge softmax stats of each step in Attention with context parallelism""" max_scale = torch.max(softmax_lse, softmax_lse_per_step) min_scale = torch.min(softmax_lse, softmax_lse_per_step) lse_diff = min_scale - max_scale - lse_diff = lse_diff.nan_to_num(nan=0.) # handle cases: tensor(-inf) - tensor(-inf) = tensor(nan); In the current cp implementation, it is possible that lses of 2 cp ranks are both -inf, if no block is selected from both cp ranks. In such cases, the finally corrected lse should remain -inf. - new_scale = max_scale + torch.log1p(torch.exp(lse_diff)) # a + ln(1 + e^(b - a)) = ln(e^a) + ln(1 + e^(b - a)) = ln(e^a + e^b) + lse_diff = lse_diff.nan_to_num( + nan=0. + ) # handle cases: tensor(-inf) - tensor(-inf) = tensor(nan); In the current cp implementation, it is possible that lses of 2 cp ranks are both -inf, if no block is selected from both cp ranks. In such cases, the finally corrected lse should remain -inf. + new_scale = max_scale + torch.log1p( + torch.exp(lse_diff)) # a + ln(1 + e^(b - a)) = ln(e^a) + ln(1 + e^(b - a)) = ln(e^a + e^b) softmax_lse.copy_(new_scale) + @torch.compile def flash_attn_fwd_out_correction_init( - out_init_step: torch.Tensor, # b h s d - softmax_lse: torch.Tensor, # b h s + out_init_step: torch.Tensor, # b h s d + softmax_lse: torch.Tensor, # b h s softmax_lse_init_step: torch.Tensor, ): """Merge partial outputs of the first step in Attention with context parallelism""" @@ -519,12 +616,14 @@ def flash_attn_fwd_out_correction( out_corrected = out_per_step * softmax_lse_corrected_exp out.add_(out_corrected) + @torch.compile def topk_sort(score, num_chunks_selected): block_indices = torch.topk(score, num_chunks_selected)[1] block_indices, _ = torch.sort(block_indices, dim=-1) return block_indices + class _attention_bsa(torch.autograd.Function): @staticmethod @@ -535,7 +634,7 @@ class _attention_bsa(torch.autograd.Function): HEAD_DIM_V = v.shape[-1] assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V assert HEAD_DIM_K in {16, 32, 64, 128, 256} - + # ---------------------- gating ---------------------- q_cmp = mean_pooling_compression(q, chunk_size_q) k_cmp = mean_pooling_compression(k, chunk_size_k) @@ -543,12 +642,8 @@ class _attention_bsa(torch.autograd.Function): # ---------------------- bsa ---------------------- - o, lse = attn_fwd_bsa_varlen_triton( - q, k, v, - sm_scale, block_indices, block_indices_lens, - chunk_size_q, chunk_size_k, - sparsity - ) + o, lse = attn_fwd_bsa_varlen_triton(q, k, v, sm_scale, block_indices, block_indices_lens, chunk_size_q, + chunk_size_k, sparsity) ctx.save_for_backward(q, k, v, o, lse, block_indices, block_indices_lens) ctx.sm_scale = sm_scale @@ -568,70 +663,59 @@ class _attention_bsa(torch.autograd.Function): dk = torch.empty_like(k) dv = torch.empty_like(v) - attn_bwd_bsa_varlen_triton( - do, - q, - k, - v, - o, - dq, - dk, - dv, - ctx.sm_scale, - lse, - block_indices, - block_indices_lens, - ctx.chunk_size_q, - ctx.chunk_size_k, - ctx.sparsity - ) + attn_bwd_bsa_varlen_triton(do, q, k, v, o, dq, dk, dv, ctx.sm_scale, lse, block_indices, block_indices_lens, + ctx.chunk_size_q, ctx.chunk_size_k, ctx.sparsity) return dq, dk, dv, None, None, None, None, None, None + flash_attn_bsa = _attention_bsa.apply + def rearrange_THW_to_3d_block(x, Nt, Nh, Nw, t, h, w, D): B, H, _, D = x.shape x = x.view(B, H, Nt, t, Nh, h, Nw, w, D) x = x.permute(0, 1, 2, 4, 6, 3, 5, 7, 8) # B H Nt Nh Nw t h w D return x.contiguous().view(B, H, Nt * Nh * Nw * t * h * w, D) + def rearrange_3d_block_to_THW(x, Nt, Nh, Nw, t, h, w, D): B, H, _, D = x.shape x = x.view(B, H, Nt, Nh, Nw, t, h, w, D) x = x.permute(0, 1, 2, 5, 3, 6, 4, 7, 8) # B H Nt t Nh h Nw w D return x.contiguous().view(B, H, Nt * t * Nh * h * Nw * w, D) + def flash_attn_bsa_3d( - q: torch.Tensor, # [B, H, Sq, D] - k: torch.Tensor, # [B, H, Skv, D] - v: torch.Tensor, # [B, H, Skv, D] + q: torch.Tensor, # [B, H, Sq, D] + k: torch.Tensor, # [B, H, Skv, D] + v: torch.Tensor, # [B, H, Skv, D] latent_shape_q, latent_shape_k, # bsa_params sparsity=0.875, cdf_threshold=None, chunk_3d_shape_q=[4, 4, 8], - chunk_3d_shape_k=[4, 4, 8], + chunk_3d_shape_k=[4, 4, 8], ) -> torch.Tensor: _, _, Sq, head_dim_q = q.shape _, _, Sk, head_dim_k = k.shape - + assert head_dim_q == head_dim_k head_dim = head_dim_q - + Tq, Hq, Wq = latent_shape_q Tk, Hk, Wk = latent_shape_k - + assert Tq * Hq * Wq == Sq assert Tk * Hk * Wk == Sk - + tq, hq, wq = chunk_3d_shape_q tk, hk, wk = chunk_3d_shape_k - + assert Tq % tq == 0 and Hq % hq == 0 and Wq % wq == 0 assert Tk % tk == 0 and Hk % hk == 0 and Wk % wk == 0 - + Ntq = Tq // tq Nhq = Hq // hq Nwq = Wq // wq @@ -650,4 +734,4 @@ def flash_attn_bsa_3d( output = flash_attn_bsa(q, k, v, chunk_size_q, chunk_size_k, sparsity, cdf_threshold, 1 / head_dim**0.5) output = rearrange_3d_block_to_THW(output, Ntq, Nhq, Nwq, tq, hq, wq, output.shape[-1]) - return output \ No newline at end of file + return output diff --git a/fastvideo/third_party/longcat_video/block_sparse_attention/common.py b/fastvideo/third_party/longcat_video/block_sparse_attention/common.py index ac817e156..09502b67f 100644 --- a/fastvideo/third_party/longcat_video/block_sparse_attention/common.py +++ b/fastvideo/third_party/longcat_video/block_sparse_attention/common.py @@ -5,11 +5,15 @@ import os if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1': autotune = triton.autotune else: + def autotune(*args, **kwargs): + def decorator(func): return func + return decorator + configs_gating_preset = { 'default': { 'BLOCK_M': 64, @@ -28,19 +32,34 @@ configs_gating = [ ] gating_reevaluate_keys = ["M", "N"] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else [] + + @autotune(configs_gating, key=gating_reevaluate_keys) @triton.jit def _attn_fwd_gating( - Q, K, Out, - stride_qz, stride_qh, stride_qm, stride_qk, - stride_kz, stride_kh, stride_kn, stride_kk, - stride_oz, stride_oh, stride_om, stride_on, - H, M, N, - HEAD_DIM: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, - ): - + Q, + K, + Out, + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_oz, + stride_oh, + stride_om, + stride_on, + H, + M, + N, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + tl.static_assert(BLOCK_N <= HEAD_DIM) start_m = tl.program_id(0) off_hz = tl.program_id(1) @@ -78,28 +97,28 @@ def _attn_fwd_gating( ) # load q: it will stay in SRAM throughout - q = tl.load(Q_block_ptr, boundary_check=(0,)) + q = tl.load(Q_block_ptr, boundary_check=(0, )) for start_n in range(0, N, BLOCK_N): start_n = tl.multiple_of(start_n, BLOCK_N) # -- compute qk ---- - k = tl.load(K_block_ptr, boundary_check=(1,)) + k = tl.load(K_block_ptr, boundary_check=(1, )) qk = tl.dot(q, k) tl.store(O_block_ptr, qk.to(Out.type.element_ty), boundary_check=(0, 1)) - + K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N)) O_block_ptr = tl.advance(O_block_ptr, (0, BLOCK_N)) @triton.jit def _attn_bwd_preprocess( - O, DO, - Delta, # output - N_CTX, - BLOCK_M: tl.constexpr, - HEAD_DIM: tl.constexpr - ): - + O, + DO, + Delta, # output + N_CTX, + BLOCK_M: tl.constexpr, + HEAD_DIM: tl.constexpr): + off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) off_hz = tl.program_id(1) off_n = tl.arange(0, HEAD_DIM) diff --git a/fastvideo/third_party/longcat_video/block_sparse_attention/communicate.py b/fastvideo/third_party/longcat_video/block_sparse_attention/communicate.py index e30524dc7..afbab5e1c 100644 --- a/fastvideo/third_party/longcat_video/block_sparse_attention/communicate.py +++ b/fastvideo/third_party/longcat_video/block_sparse_attention/communicate.py @@ -1,29 +1,19 @@ import torch -def p2p_communicate( - rank, send_tensor, send_dst, recv_tensor, recv_src, cp_group, batch_p2p_comm -): +def p2p_communicate(rank, send_tensor, send_dst, recv_tensor, recv_src, cp_group, batch_p2p_comm): """Point-to-point communications of KV and dKV in Attention with context parallelism""" send_recv_ops = [] - if batch_p2p_comm: # int(os.getenv("NVTE_BATCH_MHA_P2P_COMM", "0")) or (cp_size == 2) 为啥呢 + if batch_p2p_comm: # int(os.getenv("NVTE_BATCH_MHA_P2P_COMM", "0")) or (cp_size == 2) 为啥呢 if rank % 2 == 0: - send_op = torch.distributed.P2POp( - torch.distributed.isend, send_tensor, send_dst, cp_group - ) - recv_op = torch.distributed.P2POp( - torch.distributed.irecv, recv_tensor, recv_src, cp_group - ) + send_op = torch.distributed.P2POp(torch.distributed.isend, send_tensor, send_dst, cp_group) + recv_op = torch.distributed.P2POp(torch.distributed.irecv, recv_tensor, recv_src, cp_group) send_recv_ops.append(send_op) send_recv_ops.append(recv_op) else: - recv_op = torch.distributed.P2POp( - torch.distributed.irecv, recv_tensor, recv_src, cp_group - ) - send_op = torch.distributed.P2POp( - torch.distributed.isend, send_tensor, send_dst, cp_group - ) + recv_op = torch.distributed.P2POp(torch.distributed.irecv, recv_tensor, recv_src, cp_group) + send_op = torch.distributed.P2POp(torch.distributed.isend, send_tensor, send_dst, cp_group) send_recv_ops.append(recv_op) send_recv_ops.append(send_op) send_recv_reqs = torch.distributed.batch_isend_irecv(send_recv_ops) diff --git a/fastvideo/third_party/longcat_video/block_sparse_attention/flash_attn_bsa_varlen_mask.py b/fastvideo/third_party/longcat_video/block_sparse_attention/flash_attn_bsa_varlen_mask.py index 969cc2c39..ea52e096f 100644 --- a/fastvideo/third_party/longcat_video/block_sparse_attention/flash_attn_bsa_varlen_mask.py +++ b/fastvideo/third_party/longcat_video/block_sparse_attention/flash_attn_bsa_varlen_mask.py @@ -31,32 +31,58 @@ configs_fwd_bsa_varlen = [ for w in [4, 8] \ ] -fwd_bsa_reevaluate_varlen_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else [] +fwd_bsa_reevaluate_varlen_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_LG', 'SPARSITY'] if os.environ.get( + 'TRITON_REEVALUATE_KEY', '0') == '1' else [] + + @autotune(list(configs_fwd_bsa_varlen), key=fwd_bsa_reevaluate_varlen_keys) @triton.jit def _attn_fwd_bsa_varlen( - Q, K, V, sm_scale, M, Out, - block_indices, # [B, H, M_COMPRESS, S_MAX] - block_indices_lens, # [B, H, M_COMPRESS] - stride_qz, stride_qh, stride_qm, stride_qk, - stride_kz, stride_kh, stride_kn, stride_kk, - stride_vz, stride_vh, stride_vn, stride_vk, - stride_oz, stride_oh, stride_om, stride_ok, - stride_bz, stride_bh, stride_bm, stride_bs, - stride_lz, stride_lh, stride_lm, - H, N_CTX, - HEAD_DIM: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_N_LG: tl.constexpr, - BLOCK_N: tl.constexpr, - SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking - ): - + Q, + K, + V, + sm_scale, + M, + Out, + block_indices, # [B, H, M_COMPRESS, S_MAX] + block_indices_lens, # [B, H, M_COMPRESS] + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vn, + stride_vk, + stride_oz, + stride_oh, + stride_om, + stride_ok, + stride_bz, + stride_bh, + stride_bm, + stride_bs, + stride_lz, + stride_lh, + stride_lm, + H, + N_CTX, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N_LG: tl.constexpr, + BLOCK_N: tl.constexpr, + SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking +): + start_m = tl.program_id(0) off_hz = tl.program_id(1) off_z = off_hz // H off_h = off_hz % H - + q_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh k_offset = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh v_offset = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh @@ -129,7 +155,7 @@ def _attn_fwd_bsa_varlen( m_ij = tl.maximum(m_i, tl.max(qkT, 1) * qk_scale) qkT = qkT * qk_scale - m_ij[:, None] p = tl.math.exp2(qkT) - + # -- update m_i and l_i alpha = tl.math.exp2(m_i - m_ij) l_ij = tl.sum(p, 1) @@ -144,7 +170,7 @@ def _attn_fwd_bsa_varlen( m_i = m_ij V_block_ptr_i = tl.advance(V_block_ptr_i, (BLOCK_N, 0)) KT_block_ptr_i = tl.advance(KT_block_ptr_i, (0, BLOCK_N)) - + # epilogue m_i += tl.math.log2(l_i) acc = acc / l_i[:, None] @@ -152,6 +178,7 @@ def _attn_fwd_bsa_varlen( tl.store(m_ptrs, m_i) tl.store(O_block_ptr, acc.to(Out.type.element_ty)) + configs_fwd_bsa_varlen_align_preset = { 'default': { 'num_stages': 3, @@ -168,31 +195,57 @@ configs_fwd_bsa_varlen_align = [ for w in [4, 8] \ ] -fwd_bsa_reevaluate_varlen_align_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else [] +fwd_bsa_reevaluate_varlen_align_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_LG', 'SPARSITY'] if os.environ.get( + 'TRITON_REEVALUATE_KEY', '0') == '1' else [] + + @autotune(list(configs_fwd_bsa_varlen_align), key=fwd_bsa_reevaluate_varlen_align_keys) @triton.jit def _attn_fwd_bsa_varlen_align( - Q, K, V, sm_scale, M, Out, - block_indices, # [B, H, M_COMPRESS, S_MAX] - block_indices_lens, # [B, H, M_COMPRESS] - stride_qz, stride_qh, stride_qm, stride_qk, - stride_kz, stride_kh, stride_kn, stride_kk, - stride_vz, stride_vh, stride_vn, stride_vk, - stride_oz, stride_oh, stride_om, stride_on, - stride_bz, stride_bh, stride_bm, stride_bs, - stride_lz, stride_lh, stride_lm, - H, N_CTX, - HEAD_DIM: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_N_LG: tl.constexpr, - SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking - ): - + Q, + K, + V, + sm_scale, + M, + Out, + block_indices, # [B, H, M_COMPRESS, S_MAX] + block_indices_lens, # [B, H, M_COMPRESS] + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vn, + stride_vk, + stride_oz, + stride_oh, + stride_om, + stride_on, + stride_bz, + stride_bh, + stride_bm, + stride_bs, + stride_lz, + stride_lh, + stride_lm, + H, + N_CTX, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N_LG: tl.constexpr, + SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking +): + start_m = tl.program_id(0) off_hz = tl.program_id(1) off_z = off_hz // H off_h = off_hz % H - + q_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh k_offset = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh v_offset = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh @@ -262,7 +315,7 @@ def _attn_fwd_bsa_varlen_align( m_ij = tl.maximum(m_i, tl.max(qkT, 1) * qk_scale) qkT = qkT * qk_scale - m_ij[:, None] p = tl.math.exp2(qkT) - + # -- update m_i and l_i alpha = tl.math.exp2(m_i - m_ij) l_ij = tl.sum(p, 1) @@ -270,13 +323,12 @@ def _attn_fwd_bsa_varlen_align( acc = acc * alpha[:, None] # update acc v = tl.load(V_block_ptr_i) - acc = tl.dot(p.to(v.dtype), v, acc) # 没除se，fa2引入的优化 + acc = tl.dot(p.to(v.dtype), v, acc) # 没除se，fa2引入的优化 # update m_i and l_i # place this at the end of the loop to reduce register pressure: https://github.com/triton-lang/triton/commit/ee6abd9 - l_i = l_i * alpha + l_ij # 当前总se + l_i = l_i * alpha + l_ij # 当前总se m_i = m_ij - # epilogue m_i += tl.math.log2(l_i) acc = acc / l_i[:, None] @@ -288,24 +340,30 @@ def _attn_fwd_bsa_varlen_align( # The main inner-loop logic for computing dK and dV. @triton.jit def _attn_bwd_dkdv_bsa_varlen( - dk, dv, - k, v, - Q, DO, - M, D, + dk, + dv, + k, + v, + Q, + DO, + M, + D, block_indices, block_indices_lens, # shared by Q/K/V/DO. # stride_tok, stride_d, - stride_qm, stride_qk, - stride_dom, stride_dok, + stride_qm, + stride_qk, + stride_dom, + stride_dok, stride_mm, stride_dm, stride_bm, - N_CTX, + N_CTX, BLOCK_M1: tl.constexpr, HEAD_DIM: tl.constexpr, - ): - +): + QT_block_ptr = tl.make_block_ptr( base=Q, shape=(HEAD_DIM, N_CTX), @@ -314,7 +372,7 @@ def _attn_bwd_dkdv_bsa_varlen( block_shape=(HEAD_DIM, BLOCK_M1), order=(0, 1), ) - + DO_block_ptr = tl.make_block_ptr( base=DO, shape=(N_CTX, HEAD_DIM), @@ -329,10 +387,10 @@ def _attn_bwd_dkdv_bsa_varlen( block_id = tl.load(block_indices + i * stride_bm).to(tl.int32) start_m = block_id * BLOCK_M1 start_m = tl.multiple_of(start_m, BLOCK_M1) - + QT_block_ptr_i = tl.advance(QT_block_ptr, (0, start_m)) DO_block_ptr_i = tl.advance(DO_block_ptr, (start_m, 0)) - + qT = tl.load(QT_block_ptr_i) # Load m before computing qk to reduce pipeline stall. offs_m = start_m + tl.arange(0, BLOCK_M1) * stride_mm @@ -353,7 +411,6 @@ def _attn_bwd_dkdv_bsa_varlen( dsT = pT * (dpT - Di[None, :]) dsT = dsT.to(v.dtype) dk += tl.dot(dsT, tl.trans(qT)) - return dk, dv @@ -362,9 +419,12 @@ def _attn_bwd_dkdv_bsa_varlen( @triton.jit def _attn_bwd_dq_bsa_varlen( dq, - q, do, - m, d, - K, V, + q, + do, + m, + d, + K, + V, N_CTX, BLOCK_N2: tl.constexpr, BLOCK_N_LG: tl.constexpr, @@ -373,10 +433,12 @@ def _attn_bwd_dq_bsa_varlen( block_indices_lens, stride_bn, # stride_tok, stride_d, - stride_kn, stride_kk, - stride_vn, stride_vk, - ): - + stride_kn, + stride_kk, + stride_vn, + stride_vk, +): + VT_block_ptr = tl.make_block_ptr( base=V, shape=(HEAD_DIM, N_CTX), @@ -393,19 +455,19 @@ def _attn_bwd_dq_bsa_varlen( block_shape=(HEAD_DIM, BLOCK_N2), order=(0, 1), ) - + S = tl.load(block_indices_lens) for i in range(S): block_id = tl.load(block_indices + i * stride_bn).to(tl.int32) lo, hi = block_id * BLOCK_N_LG, (block_id + 1) * BLOCK_N_LG lo = tl.multiple_of(lo, BLOCK_N2) - + KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo)) VT_block_ptr_i = tl.advance(VT_block_ptr, (0, lo)) - + for start_n in range(lo, hi, BLOCK_N2): start_n = tl.multiple_of(start_n, BLOCK_N2) - + kT = tl.load(KT_block_ptr_i) vT = tl.load(VT_block_ptr_i) qkT = tl.dot(q, kT) @@ -413,33 +475,41 @@ def _attn_bwd_dq_bsa_varlen( # Compute dP and dS. dp = tl.dot(do, vT).to(tl.float32) ds = p * (dp - d) - ds = ds.to(kT.dtype) # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py: Converting ds to q.dtype here reduces register pressure and makes it much faster for BLOCK_HEADDIM=128 + ds = ds.to( + kT.dtype + ) # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py: Converting ds to q.dtype here reduces register pressure and makes it much faster for BLOCK_HEADDIM=128 # Compute dQ. # NOTE: We need to de-scale dq in the end, because kT was pre-scaled. dq += tl.dot(ds, tl.trans(kT)) # Increment pointers. KT_block_ptr_i = tl.advance(KT_block_ptr_i, (0, BLOCK_N2)) VT_block_ptr_i = tl.advance(VT_block_ptr_i, (0, BLOCK_N2)) - + return dq + # the main inner-loop logic for computing dQ @triton.jit def _attn_bwd_dq_bsa_varlen_align( dq, - q, do, - m, d, - K, V, + q, + do, + m, + d, + K, + V, N_CTX, BLOCK_N_LG: tl.constexpr, HEAD_DIM: tl.constexpr, block_indices, block_indices_lens, stride_bn, - stride_kn, stride_kk, - stride_vn, stride_vk, - ): - + stride_kn, + stride_kk, + stride_vn, + stride_vk, +): + VT_block_ptr = tl.make_block_ptr( base=V, shape=(HEAD_DIM, N_CTX), @@ -456,16 +526,16 @@ def _attn_bwd_dq_bsa_varlen_align( block_shape=(HEAD_DIM, BLOCK_N_LG), order=(0, 1), ) - + S = tl.load(block_indices_lens) for i in range(S): block_id = tl.load(block_indices + i * stride_bn).to(tl.int32) lo = block_id * BLOCK_N_LG lo = tl.multiple_of(lo, BLOCK_N_LG) - + KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo)) VT_block_ptr_i = tl.advance(VT_block_ptr, (0, lo)) - + kT = tl.load(KT_block_ptr_i) vT = tl.load(VT_block_ptr_i) qkT = tl.dot(q, kT) @@ -473,14 +543,16 @@ def _attn_bwd_dq_bsa_varlen_align( # Compute dP and dS. dp = tl.dot(do, vT).to(tl.float32) ds = p * (dp - d) - ds = ds.to(kT.dtype) # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py: Converting ds to q.dtype here reduces register pressure and makes it much faster for BLOCK_HEADDIM=128 + ds = ds.to( + kT.dtype + ) # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py: Converting ds to q.dtype here reduces register pressure and makes it much faster for BLOCK_HEADDIM=128 # Compute dQ. # NOTE: We need to de-scale dq in the end, because kT was pre-scaled. dq += tl.dot(ds, tl.trans(kT)) - return dq + configs_bwd_dkdv_bsa_varlen_preset = { 'default': { 'BLOCK_N': 128, @@ -499,40 +571,75 @@ configs_bwd_dkdv_bsa_varlen = [ for s in [2, 3, 4, 5] \ for w in [4, 8] \ ] -bwd_dkdv_bsa_varlen_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else [] +bwd_dkdv_bsa_varlen_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get( + 'TRITON_REEVALUATE_KEY', '0') == '1' else [] + + @autotune(list(configs_bwd_dkdv_bsa_varlen), key=bwd_dkdv_bsa_varlen_reevaluate_keys) @triton.jit def _attn_bwd_dkdv_bsa_varlen_wrapper( - Q, K, V, sm_scale, # softmax scale - DO, - DK, DV, - M, # lse (log2) - D, - block_indices, - block_indices_lens, - # stride_z, stride_h, stride_tok, stride_d, # shared by Q/K/V/DO. - # qkv - stride_qz, stride_qh, stride_qm, stride_qk, - stride_kz, stride_kh, stride_kn, stride_kk, - stride_vz, stride_vh, stride_vn, stride_vk, - # dk dv do - stride_dkz, stride_dkh, stride_dkn, stride_dkk, - stride_dvz, stride_dvh, stride_dvn, stride_dvk, - stride_doz, stride_doh, stride_dom, stride_dok, - # m, d - stride_mz, stride_mh, stride_mm, - stride_dz, stride_dh, stride_dm, - # - stride_bz, stride_bh, stride_bn, stride_bm, # block_indices - stride_lz, stride_lh, stride_ln, # block_indices_lens - # - H, N_CTX, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_N_DQ_LG: tl.constexpr, # logical block size - HEAD_DIM: tl.constexpr, - SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking - ): + Q, + K, + V, + sm_scale, # softmax scale + DO, + DK, + DV, + M, # lse (log2) + D, + block_indices, + block_indices_lens, + # stride_z, stride_h, stride_tok, stride_d, # shared by Q/K/V/DO. + # qkv + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vn, + stride_vk, + # dk dv do + stride_dkz, + stride_dkh, + stride_dkn, + stride_dkk, + stride_dvz, + stride_dvh, + stride_dvn, + stride_dvk, + stride_doz, + stride_doh, + stride_dom, + stride_dok, + # m, d + stride_mz, + stride_mh, + stride_mm, + stride_dz, + stride_dh, + stride_dm, + # + stride_bz, + stride_bh, + stride_bn, + stride_bm, # block_indices + stride_lz, + stride_lh, + stride_ln, # block_indices_lens + # + H, + N_CTX, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_N_DQ_LG: tl.constexpr, # logical block size + HEAD_DIM: tl.constexpr, + SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking +): tl.static_assert(BLOCK_N_DQ_LG % BLOCK_N == 0) start_n = tl.program_id(0) @@ -540,7 +647,7 @@ def _attn_bwd_dkdv_bsa_varlen_wrapper( off_hz = tl.program_id(2) off_z = off_hz // H off_h = off_hz % H - + off_q = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh off_k = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh off_v = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh @@ -548,13 +655,13 @@ def _attn_bwd_dkdv_bsa_varlen_wrapper( off_dk = off_z.to(tl.int64) * stride_dkz + off_h.to(tl.int64) * stride_dkh off_dv = off_z.to(tl.int64) * stride_dvz + off_h.to(tl.int64) * stride_dvh off_do = off_z.to(tl.int64) * stride_doz + off_h.to(tl.int64) * stride_doh - + off_m = off_z.to(tl.int64) * stride_mz + off_h.to(tl.int64) * stride_mh off_d = off_z.to(tl.int64) * stride_dz + off_h.to(tl.int64) * stride_dh - off_block_incides = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh + off_block_incides = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh off_block_incides_lens = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh - + # offset pointers for batch/head Q += off_q K += off_k @@ -562,7 +669,7 @@ def _attn_bwd_dkdv_bsa_varlen_wrapper( DO += off_do DK += off_dk DV += off_dv - + M += off_m D += off_d block_indices += off_block_incides @@ -574,7 +681,7 @@ def _attn_bwd_dkdv_bsa_varlen_wrapper( dk = tl.zeros([BLOCK_N, HEAD_DIM], dtype=tl.float32) # load K and V: they stay in SRAM throughout the inner loop. - + K_block_ptr = tl.make_block_ptr( base=K, shape=(N_CTX, HEAD_DIM), @@ -607,40 +714,47 @@ def _attn_bwd_dkdv_bsa_varlen_wrapper( block_shape=(BLOCK_N, HEAD_DIM), order=(1, 0), ) - + k = tl.load(K_block_ptr) v = tl.load(V_block_ptr) - + k_compress_idx = start_n * BLOCK_N // BLOCK_N_DQ_LG block_indices_i = block_indices + k_compress_idx * stride_bn block_indices_lens_i = block_indices_lens + k_compress_idx * stride_ln - + dk, dv = _attn_bwd_dkdv_bsa_varlen( - dk, dv, - k, v, - Q, DO, - M, D, + dk, + dv, + k, + v, + Q, + DO, + M, + D, block_indices_i, block_indices_lens_i, # shared by Q/K/V/DO. - stride_qm, stride_qk, - stride_dom, stride_dok, + stride_qm, + stride_qk, + stride_dom, + stride_dok, stride_mm, stride_dm, # stride_bm, - N_CTX, - BLOCK_M, - HEAD_DIM, + N_CTX, + BLOCK_M, + HEAD_DIM, ) - + # Write back dk - dk *= sm_scale # S = scale * QKT; dK = scale * QdST + dk *= sm_scale # S = scale * QKT; dK = scale * QdST tl.store(DK_block_ptr, dk.to(k.dtype)) - + # Write back dv tl.store(DV_block_ptr, dv.to(v.dtype)) + configs_bwd_dq_bsa_varlen_preset = { 'default': { 'BLOCK_N_DQ': 64, @@ -659,62 +773,92 @@ configs_bwd_dq_bsa_varlen = [ for s in [2, 3, 4, 5] \ for w in [4, 8] \ ] -bwd_dq_bsa_varlen_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else [] +bwd_dq_bsa_varlen_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get( + 'TRITON_REEVALUATE_KEY', '0') == '1' else [] + + @autotune(list(configs_bwd_dq_bsa_varlen), key=bwd_dq_bsa_varlen_reevaluate_keys) @triton.jit def _attn_bwd_dq_bsa_varlen_wrapper( - Q, K, V, # softmax scale - DO, - DQ, - M, # lse (log2) - D, - block_indices, - block_indices_lens, - # stride_z, stride_h, stride_tok, stride_d, # shared by Q/K/V/DO. - # qkv - stride_qz, stride_qh, stride_qm, stride_qk, - stride_kz, stride_kh, stride_kn, stride_kk, - stride_vz, stride_vh, stride_vn, stride_vk, - # dq do - stride_dqz, stride_dqh, stride_dqm, stride_dqk, - stride_doz, stride_doh, stride_dom, stride_dok, - # m, d - stride_mz, stride_mh, stride_mm, - stride_dz, stride_dh, stride_dm, - # - stride_bz, stride_bh, stride_bm, stride_bn, # block_indices - stride_lz, stride_lh, stride_lm, # block_indices_lens - # - H, N_CTX, - BLOCK_M: tl.constexpr, - BLOCK_N_DQ_LG: tl.constexpr, # logical block size - BLOCK_N_DQ: tl.constexpr, - HEAD_DIM: tl.constexpr, - SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking - ): + Q, + K, + V, # softmax scale + DO, + DQ, + M, # lse (log2) + D, + block_indices, + block_indices_lens, + # stride_z, stride_h, stride_tok, stride_d, # shared by Q/K/V/DO. + # qkv + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vn, + stride_vk, + # dq do + stride_dqz, + stride_dqh, + stride_dqm, + stride_dqk, + stride_doz, + stride_doh, + stride_dom, + stride_dok, + # m, d + stride_mz, + stride_mh, + stride_mm, + stride_dz, + stride_dh, + stride_dm, + # + stride_bz, + stride_bh, + stride_bm, + stride_bn, # block_indices + stride_lz, + stride_lh, + stride_lm, # block_indices_lens + # + H, + N_CTX, + BLOCK_M: tl.constexpr, + BLOCK_N_DQ_LG: tl.constexpr, # logical block size + BLOCK_N_DQ: tl.constexpr, + HEAD_DIM: tl.constexpr, + SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking +): tl.static_assert(BLOCK_N_DQ_LG % BLOCK_N_DQ == 0) tl.static_assert(BLOCK_N_DQ_LG % BLOCK_M == 0) - + LN2: tl.constexpr = 0.6931471824645996 # = ln(2) - + start_m = tl.program_id(0) off_hz = tl.program_id(2) off_z = off_hz // H off_h = off_hz % H - + off_q = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh off_k = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh off_v = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh off_dq = off_z.to(tl.int64) * stride_dqz + off_h.to(tl.int64) * stride_dqh off_do = off_z.to(tl.int64) * stride_doz + off_h.to(tl.int64) * stride_doh - + off_m = off_z.to(tl.int64) * stride_mz + off_h.to(tl.int64) * stride_mh off_d = off_z.to(tl.int64) * stride_dz + off_h.to(tl.int64) * stride_dh - off_block_incides = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh + off_block_incides = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh off_block_incides_lens = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh # offset pointers for batch/head @@ -723,14 +867,14 @@ def _attn_bwd_dq_bsa_varlen_wrapper( V += off_v DO += off_do DQ += off_dq - + M += off_m D += off_d block_indices += off_block_incides block_indices_lens += off_block_incides_lens # ---------------------------------------- [DQ] ---------------------------------------- - + Q_block_ptr = tl.make_block_ptr( base=Q, shape=(N_CTX, HEAD_DIM), @@ -755,12 +899,12 @@ def _attn_bwd_dq_bsa_varlen_wrapper( block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0), ) - + q = tl.load(Q_block_ptr) do = tl.load(DO_block_ptr) - + dq = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) - + start_m = start_m * BLOCK_M offs_m = start_m + tl.arange(0, BLOCK_M) * stride_mm @@ -771,30 +915,36 @@ def _attn_bwd_dq_bsa_varlen_wrapper( d = tl.load(D + offs_d) d = d[:, None] - + block_indices_m = block_indices + (start_m // BLOCK_M) * stride_bm block_indices_lens_m = block_indices_lens + (start_m // BLOCK_M) * stride_lm dq = _attn_bwd_dq_bsa_varlen( dq, - q, do, - m, d, - K, V, + q, + do, + m, + d, + K, + V, N_CTX, BLOCK_N_DQ, - BLOCK_N_DQ_LG, + BLOCK_N_DQ_LG, HEAD_DIM, - block_indices_m, - block_indices_lens_m, - stride_bn, - stride_kn, stride_kk, - stride_vn, stride_vk, + block_indices_m, + block_indices_lens_m, + stride_bn, + stride_kn, + stride_kk, + stride_vn, + stride_vk, ) - + # Write back dQ. dq *= LN2 tl.store(DQ_block_ptr, dq.to(q.dtype)) + configs_bwd_dq_bsa_varlen_align_preset = { 'default': { 'num_stages': 2, @@ -810,60 +960,90 @@ configs_bwd_dq_bsa_varlen_align = [ for s in [2, 3, 4, 5] \ for w in [4, 8] \ ] -bwd_dq_bsa_varlen_align_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else [] +bwd_dq_bsa_varlen_align_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get( + 'TRITON_REEVALUATE_KEY', '0') == '1' else [] + + @autotune(list(configs_bwd_dq_bsa_varlen_align), key=bwd_dq_bsa_varlen_align_reevaluate_keys) @triton.jit def _attn_bwd_dq_bsa_varlen_align_wrapper( - Q, K, V, # softmax scale - DO, - DQ, - M, # lse (log2) - D, - block_indices, - block_indices_lens, - # qkv - stride_qz, stride_qh, stride_qm, stride_qk, - stride_kz, stride_kh, stride_kn, stride_kk, - stride_vz, stride_vh, stride_vn, stride_vk, - # dq do - stride_dqz, stride_dqh, stride_dqm, stride_dqk, - stride_doz, stride_doh, stride_dom, stride_dok, - # m, d - stride_mz, stride_mh, stride_mm, - stride_dz, stride_dh, stride_dm, - # - stride_bz, stride_bh, stride_bm, stride_bn, # block_indices - stride_lz, stride_lh, stride_lm, # block_indices_lens - # - H, N_CTX, - BLOCK_M: tl.constexpr, - BLOCK_N_DQ_LG: tl.constexpr, # logical block size - HEAD_DIM: tl.constexpr, - SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking - ): + Q, + K, + V, # softmax scale + DO, + DQ, + M, # lse (log2) + D, + block_indices, + block_indices_lens, + # qkv + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vn, + stride_vk, + # dq do + stride_dqz, + stride_dqh, + stride_dqm, + stride_dqk, + stride_doz, + stride_doh, + stride_dom, + stride_dok, + # m, d + stride_mz, + stride_mh, + stride_mm, + stride_dz, + stride_dh, + stride_dm, + # + stride_bz, + stride_bh, + stride_bm, + stride_bn, # block_indices + stride_lz, + stride_lh, + stride_lm, # block_indices_lens + # + H, + N_CTX, + BLOCK_M: tl.constexpr, + BLOCK_N_DQ_LG: tl.constexpr, # logical block size + HEAD_DIM: tl.constexpr, + SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking +): tl.static_assert(BLOCK_N_DQ_LG % BLOCK_N_DQ_LG == 0) tl.static_assert(BLOCK_N_DQ_LG % BLOCK_M == 0) - + LN2: tl.constexpr = 0.6931471824645996 # = ln(2) - + start_m = tl.program_id(0) - + off_hz = tl.program_id(2) off_z = off_hz // H off_h = off_hz % H - + off_q = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh off_k = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh off_v = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh off_dq = off_z.to(tl.int64) * stride_dqz + off_h.to(tl.int64) * stride_dqh off_do = off_z.to(tl.int64) * stride_doz + off_h.to(tl.int64) * stride_doh - + off_m = off_z.to(tl.int64) * stride_mz + off_h.to(tl.int64) * stride_mh off_d = off_z.to(tl.int64) * stride_dz + off_h.to(tl.int64) * stride_dh - off_block_incides = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh + off_block_incides = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh off_block_incides_lens = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh # offset pointers for batch/head @@ -872,14 +1052,14 @@ def _attn_bwd_dq_bsa_varlen_align_wrapper( V += off_v DO += off_do DQ += off_dq - + M += off_m D += off_d block_indices += off_block_incides block_indices_lens += off_block_incides_lens # ---------------------------------------- [DQ] ---------------------------------------- - + Q_block_ptr = tl.make_block_ptr( base=Q, shape=(N_CTX, HEAD_DIM), @@ -904,12 +1084,12 @@ def _attn_bwd_dq_bsa_varlen_align_wrapper( block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0), ) - + q = tl.load(Q_block_ptr) do = tl.load(DO_block_ptr) - + dq = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) - + start_m = start_m * BLOCK_M offs_m = start_m + tl.arange(0, BLOCK_M) * stride_mm @@ -917,30 +1097,34 @@ def _attn_bwd_dq_bsa_varlen_align_wrapper( m = tl.load(M + offs_m) m = m[:, None] - + # D (= delta) is pre-divided by ds_scale. d = tl.load(D + offs_d) d = d[:, None] - + block_indices_m = block_indices + (start_m // BLOCK_M) * stride_bm block_indices_lens_m = block_indices_lens + (start_m // BLOCK_M) * stride_lm dq = _attn_bwd_dq_bsa_varlen_align( dq, - q, do, - m, d, - K, V, + q, + do, + m, + d, + K, + V, N_CTX, - BLOCK_N_DQ_LG, + BLOCK_N_DQ_LG, HEAD_DIM, - block_indices_m, - block_indices_lens_m, - stride_bn, - stride_kn, stride_kk, - stride_vn, stride_vk, + block_indices_m, + block_indices_lens_m, + stride_bn, + stride_kn, + stride_kk, + stride_vn, + stride_vk, ) - - # Write back dQ. + + # Write back dQ. dq *= LN2 tl.store(DQ_block_ptr, dq.to(q.dtype)) - diff --git a/fastvideo/third_party/pynvml.py b/fastvideo/third_party/pynvml.py index 0a4be23a0..2d4a09c2d 100644 --- a/fastvideo/third_party/pynvml.py +++ b/fastvideo/third_party/pynvml.py @@ -44,52 +44,51 @@ import string ## C Type mappings ## ## Enums _nvmlEnableState_t = c_uint -NVML_FEATURE_DISABLED = 0 -NVML_FEATURE_ENABLED = 1 +NVML_FEATURE_DISABLED = 0 +NVML_FEATURE_ENABLED = 1 _nvmlBrandType_t = c_uint -NVML_BRAND_UNKNOWN = 0 -NVML_BRAND_QUADRO = 1 -NVML_BRAND_TESLA = 2 -NVML_BRAND_NVS = 3 -NVML_BRAND_GRID = 4 # Deprecated from API reporting. Keeping definition for backward compatibility. -NVML_BRAND_GEFORCE = 5 -NVML_BRAND_TITAN = 6 -NVML_BRAND_NVIDIA_VAPPS = 7 # NVIDIA Virtual Applications -NVML_BRAND_NVIDIA_VPC = 8 # NVIDIA Virtual PC -NVML_BRAND_NVIDIA_VCS = 9 # NVIDIA Virtual Compute Server -NVML_BRAND_NVIDIA_VWS = 10 # NVIDIA RTX Virtual Workstation +NVML_BRAND_UNKNOWN = 0 +NVML_BRAND_QUADRO = 1 +NVML_BRAND_TESLA = 2 +NVML_BRAND_NVS = 3 +NVML_BRAND_GRID = 4 # Deprecated from API reporting. Keeping definition for backward compatibility. +NVML_BRAND_GEFORCE = 5 +NVML_BRAND_TITAN = 6 +NVML_BRAND_NVIDIA_VAPPS = 7 # NVIDIA Virtual Applications +NVML_BRAND_NVIDIA_VPC = 8 # NVIDIA Virtual PC +NVML_BRAND_NVIDIA_VCS = 9 # NVIDIA Virtual Compute Server +NVML_BRAND_NVIDIA_VWS = 10 # NVIDIA RTX Virtual Workstation NVML_BRAND_NVIDIA_CLOUD_GAMING = 11 # NVIDIA Cloud Gaming -NVML_BRAND_NVIDIA_VGAMING = NVML_BRAND_NVIDIA_CLOUD_GAMING # Deprecated from API reporting. Keeping definition for backward compatibility. -NVML_BRAND_QUADRO_RTX = 12 -NVML_BRAND_NVIDIA_RTX = 13 -NVML_BRAND_NVIDIA = 14 -NVML_BRAND_GEFORCE_RTX = 15 # Unused -NVML_BRAND_TITAN_RTX = 16 # Unused -NVML_BRAND_COUNT = 17 +NVML_BRAND_NVIDIA_VGAMING = NVML_BRAND_NVIDIA_CLOUD_GAMING # Deprecated from API reporting. Keeping definition for backward compatibility. +NVML_BRAND_QUADRO_RTX = 12 +NVML_BRAND_NVIDIA_RTX = 13 +NVML_BRAND_NVIDIA = 14 +NVML_BRAND_GEFORCE_RTX = 15 # Unused +NVML_BRAND_TITAN_RTX = 16 # Unused +NVML_BRAND_COUNT = 17 _nvmlTemperatureThresholds_t = c_uint -NVML_TEMPERATURE_THRESHOLD_SHUTDOWN = 0 -NVML_TEMPERATURE_THRESHOLD_SLOWDOWN = 1 -NVML_TEMPERATURE_THRESHOLD_MEM_MAX = 2 -NVML_TEMPERATURE_THRESHOLD_GPU_MAX = 3 -NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MIN = 4 +NVML_TEMPERATURE_THRESHOLD_SHUTDOWN = 0 +NVML_TEMPERATURE_THRESHOLD_SLOWDOWN = 1 +NVML_TEMPERATURE_THRESHOLD_MEM_MAX = 2 +NVML_TEMPERATURE_THRESHOLD_GPU_MAX = 3 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MIN = 4 NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_CURR = 5 -NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MAX = 6 -NVML_TEMPERATURE_THRESHOLD_GPS_CURR = 7 -NVML_TEMPERATURE_THRESHOLD_COUNT = 8 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MAX = 6 +NVML_TEMPERATURE_THRESHOLD_GPS_CURR = 7 +NVML_TEMPERATURE_THRESHOLD_COUNT = 8 _nvmlTemperatureSensors_t = c_uint -NVML_TEMPERATURE_GPU = 0 -NVML_TEMPERATURE_COUNT = 1 - +NVML_TEMPERATURE_GPU = 0 +NVML_TEMPERATURE_COUNT = 1 _nvmlComputeMode_t = c_uint -NVML_COMPUTEMODE_DEFAULT = 0 -NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1 ## Support Removed -NVML_COMPUTEMODE_PROHIBITED = 2 +NVML_COMPUTEMODE_DEFAULT = 0 +NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1 ## Support Removed +NVML_COMPUTEMODE_PROHIBITED = 2 NVML_COMPUTEMODE_EXCLUSIVE_PROCESS = 3 -NVML_COMPUTEMODE_COUNT = 4 +NVML_COMPUTEMODE_COUNT = 4 _nvmlMemoryLocation_t = c_uint NVML_MEMORY_LOCATION_L1_CACHE = 0 @@ -126,158 +125,158 @@ NVML_NVLINK_ERROR_DL_ECC_COUNT = 5 _nvmlNvLinkCapability_t = c_uint NVML_NVLINK_CAP_P2P_SUPPORTED = 0 NVML_NVLINK_CAP_SYSMEM_ACCESS = 1 -NVML_NVLINK_CAP_P2P_ATOMICS = 2 -NVML_NVLINK_CAP_SYSMEM_ATOMICS= 3 -NVML_NVLINK_CAP_SLI_BRIDGE = 4 -NVML_NVLINK_CAP_VALID = 5 -NVML_NVLINK_CAP_COUNT = 6 +NVML_NVLINK_CAP_P2P_ATOMICS = 2 +NVML_NVLINK_CAP_SYSMEM_ATOMICS = 3 +NVML_NVLINK_CAP_SLI_BRIDGE = 4 +NVML_NVLINK_CAP_VALID = 5 +NVML_NVLINK_CAP_COUNT = 6 _nvmlNvLinkUtilizationCountPktTypes_t = c_uint -NVML_NVLINK_COUNTER_PKTFILTER_NOP = 0x1 -NVML_NVLINK_COUNTER_PKTFILTER_READ = 0x2 -NVML_NVLINK_COUNTER_PKTFILTER_WRITE = 0x4 -NVML_NVLINK_COUNTER_PKTFILTER_RATOM = 0x8 -NVML_NVLINK_COUNTER_PKTFILTER_NRATOM = 0x10 -NVML_NVLINK_COUNTER_PKTFILTER_FLUSH = 0x20 -NVML_NVLINK_COUNTER_PKTFILTER_RESPDATA = 0x40 +NVML_NVLINK_COUNTER_PKTFILTER_NOP = 0x1 +NVML_NVLINK_COUNTER_PKTFILTER_READ = 0x2 +NVML_NVLINK_COUNTER_PKTFILTER_WRITE = 0x4 +NVML_NVLINK_COUNTER_PKTFILTER_RATOM = 0x8 +NVML_NVLINK_COUNTER_PKTFILTER_NRATOM = 0x10 +NVML_NVLINK_COUNTER_PKTFILTER_FLUSH = 0x20 +NVML_NVLINK_COUNTER_PKTFILTER_RESPDATA = 0x40 NVML_NVLINK_COUNTER_PKTFILTER_RESPNODATA = 0x80 -NVML_NVLINK_COUNTER_PKTFILTER_ALL = 0xFF +NVML_NVLINK_COUNTER_PKTFILTER_ALL = 0xFF _nvmlNvLinkUtilizationCountUnits_t = c_uint -NVML_NVLINK_COUNTER_UNIT_CYCLES = 0 -NVML_NVLINK_COUNTER_UNIT_PACKETS = 1 -NVML_NVLINK_COUNTER_UNIT_BYTES = 2 +NVML_NVLINK_COUNTER_UNIT_CYCLES = 0 +NVML_NVLINK_COUNTER_UNIT_PACKETS = 1 +NVML_NVLINK_COUNTER_UNIT_BYTES = 2 NVML_NVLINK_COUNTER_UNIT_RESERVED = 3 -NVML_NVLINK_COUNTER_UNIT_COUNT = 4 +NVML_NVLINK_COUNTER_UNIT_COUNT = 4 _nvmlNvLinkDeviceType_t = c_uint -NVML_NVLINK_DEVICE_TYPE_GPU = 0x00 -NVML_NVLINK_DEVICE_TYPE_IBMNPU = 0x01 -NVML_NVLINK_DEVICE_TYPE_SWITCH = 0x02 +NVML_NVLINK_DEVICE_TYPE_GPU = 0x00 +NVML_NVLINK_DEVICE_TYPE_IBMNPU = 0x01 +NVML_NVLINK_DEVICE_TYPE_SWITCH = 0x02 NVML_NVLINK_DEVICE_TYPE_UNKNOWN = 0xFF # These are deprecated, instead use _nvmlMemoryErrorType_t _nvmlEccBitType_t = c_uint -NVML_SINGLE_BIT_ECC = 0 -NVML_DOUBLE_BIT_ECC = 1 +NVML_SINGLE_BIT_ECC = 0 +NVML_DOUBLE_BIT_ECC = 1 NVML_ECC_ERROR_TYPE_COUNT = 2 _nvmlEccCounterType_t = c_uint -NVML_VOLATILE_ECC = 0 -NVML_AGGREGATE_ECC = 1 +NVML_VOLATILE_ECC = 0 +NVML_AGGREGATE_ECC = 1 NVML_ECC_COUNTER_TYPE_COUNT = 2 _nvmlMemoryErrorType_t = c_uint -NVML_MEMORY_ERROR_TYPE_CORRECTED = 0 +NVML_MEMORY_ERROR_TYPE_CORRECTED = 0 NVML_MEMORY_ERROR_TYPE_UNCORRECTED = 1 -NVML_MEMORY_ERROR_TYPE_COUNT = 2 +NVML_MEMORY_ERROR_TYPE_COUNT = 2 _nvmlClockType_t = c_uint -NVML_CLOCK_GRAPHICS = 0 -NVML_CLOCK_SM = 1 -NVML_CLOCK_MEM = 2 -NVML_CLOCK_VIDEO = 3 -NVML_CLOCK_COUNT = 4 +NVML_CLOCK_GRAPHICS = 0 +NVML_CLOCK_SM = 1 +NVML_CLOCK_MEM = 2 +NVML_CLOCK_VIDEO = 3 +NVML_CLOCK_COUNT = 4 _nvmlClockId_t = c_uint -NVML_CLOCK_ID_CURRENT = 0 -NVML_CLOCK_ID_APP_CLOCK_TARGET = 1 -NVML_CLOCK_ID_APP_CLOCK_DEFAULT = 2 +NVML_CLOCK_ID_CURRENT = 0 +NVML_CLOCK_ID_APP_CLOCK_TARGET = 1 +NVML_CLOCK_ID_APP_CLOCK_DEFAULT = 2 NVML_CLOCK_ID_CUSTOMER_BOOST_MAX = 3 -NVML_CLOCK_ID_COUNT = 4 +NVML_CLOCK_ID_COUNT = 4 _nvmlDriverModel_t = c_uint -NVML_DRIVER_WDDM = 0 -NVML_DRIVER_WDM = 1 -NVML_DRIVER_MCDM = 2 +NVML_DRIVER_WDDM = 0 +NVML_DRIVER_WDM = 1 +NVML_DRIVER_MCDM = 2 NVML_MAX_GPU_PERF_PSTATES = 16 _nvmlPstates_t = c_uint -NVML_PSTATE_0 = 0 -NVML_PSTATE_1 = 1 -NVML_PSTATE_2 = 2 -NVML_PSTATE_3 = 3 -NVML_PSTATE_4 = 4 -NVML_PSTATE_5 = 5 -NVML_PSTATE_6 = 6 -NVML_PSTATE_7 = 7 -NVML_PSTATE_8 = 8 -NVML_PSTATE_9 = 9 -NVML_PSTATE_10 = 10 -NVML_PSTATE_11 = 11 -NVML_PSTATE_12 = 12 -NVML_PSTATE_13 = 13 -NVML_PSTATE_14 = 14 -NVML_PSTATE_15 = 15 -NVML_PSTATE_UNKNOWN = 32 +NVML_PSTATE_0 = 0 +NVML_PSTATE_1 = 1 +NVML_PSTATE_2 = 2 +NVML_PSTATE_3 = 3 +NVML_PSTATE_4 = 4 +NVML_PSTATE_5 = 5 +NVML_PSTATE_6 = 6 +NVML_PSTATE_7 = 7 +NVML_PSTATE_8 = 8 +NVML_PSTATE_9 = 9 +NVML_PSTATE_10 = 10 +NVML_PSTATE_11 = 11 +NVML_PSTATE_12 = 12 +NVML_PSTATE_13 = 13 +NVML_PSTATE_14 = 14 +NVML_PSTATE_15 = 15 +NVML_PSTATE_UNKNOWN = 32 _nvmlInforomObject_t = c_uint -NVML_INFOROM_OEM = 0 -NVML_INFOROM_ECC = 1 -NVML_INFOROM_POWER = 2 -NVML_INFOROM_DEN = 3 -NVML_INFOROM_COUNT = 4 +NVML_INFOROM_OEM = 0 +NVML_INFOROM_ECC = 1 +NVML_INFOROM_POWER = 2 +NVML_INFOROM_DEN = 3 +NVML_INFOROM_COUNT = 4 _nvmlReturn_t = c_uint -NVML_SUCCESS = 0 -NVML_ERROR_UNINITIALIZED = 1 -NVML_ERROR_INVALID_ARGUMENT = 2 -NVML_ERROR_NOT_SUPPORTED = 3 -NVML_ERROR_NO_PERMISSION = 4 -NVML_ERROR_ALREADY_INITIALIZED = 5 -NVML_ERROR_NOT_FOUND = 6 -NVML_ERROR_INSUFFICIENT_SIZE = 7 -NVML_ERROR_INSUFFICIENT_POWER = 8 -NVML_ERROR_DRIVER_NOT_LOADED = 9 -NVML_ERROR_TIMEOUT = 10 -NVML_ERROR_IRQ_ISSUE = 11 -NVML_ERROR_LIBRARY_NOT_FOUND = 12 -NVML_ERROR_FUNCTION_NOT_FOUND = 13 -NVML_ERROR_CORRUPTED_INFOROM = 14 -NVML_ERROR_GPU_IS_LOST = 15 -NVML_ERROR_RESET_REQUIRED = 16 -NVML_ERROR_OPERATING_SYSTEM = 17 -NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18 -NVML_ERROR_IN_USE = 19 -NVML_ERROR_MEMORY = 20 -NVML_ERROR_NO_DATA = 21 -NVML_ERROR_VGPU_ECC_NOT_SUPPORTED = 22 -NVML_ERROR_INSUFFICIENT_RESOURCES = 23 -NVML_ERROR_FREQ_NOT_SUPPORTED = 24 +NVML_SUCCESS = 0 +NVML_ERROR_UNINITIALIZED = 1 +NVML_ERROR_INVALID_ARGUMENT = 2 +NVML_ERROR_NOT_SUPPORTED = 3 +NVML_ERROR_NO_PERMISSION = 4 +NVML_ERROR_ALREADY_INITIALIZED = 5 +NVML_ERROR_NOT_FOUND = 6 +NVML_ERROR_INSUFFICIENT_SIZE = 7 +NVML_ERROR_INSUFFICIENT_POWER = 8 +NVML_ERROR_DRIVER_NOT_LOADED = 9 +NVML_ERROR_TIMEOUT = 10 +NVML_ERROR_IRQ_ISSUE = 11 +NVML_ERROR_LIBRARY_NOT_FOUND = 12 +NVML_ERROR_FUNCTION_NOT_FOUND = 13 +NVML_ERROR_CORRUPTED_INFOROM = 14 +NVML_ERROR_GPU_IS_LOST = 15 +NVML_ERROR_RESET_REQUIRED = 16 +NVML_ERROR_OPERATING_SYSTEM = 17 +NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18 +NVML_ERROR_IN_USE = 19 +NVML_ERROR_MEMORY = 20 +NVML_ERROR_NO_DATA = 21 +NVML_ERROR_VGPU_ECC_NOT_SUPPORTED = 22 +NVML_ERROR_INSUFFICIENT_RESOURCES = 23 +NVML_ERROR_FREQ_NOT_SUPPORTED = 24 NVML_ERROR_ARGUMENT_VERSION_MISMATCH = 25 -NVML_ERROR_DEPRECATED = 26 -NVML_ERROR_NOT_READY = 27 -NVML_ERROR_GPU_NOT_FOUND = 28 -NVML_ERROR_INVALID_STATE = 29 -NVML_ERROR_UNKNOWN = 999 +NVML_ERROR_DEPRECATED = 26 +NVML_ERROR_NOT_READY = 27 +NVML_ERROR_GPU_NOT_FOUND = 28 +NVML_ERROR_INVALID_STATE = 29 +NVML_ERROR_UNKNOWN = 999 _nvmlFanState_t = c_uint -NVML_FAN_NORMAL = 0 -NVML_FAN_FAILED = 1 +NVML_FAN_NORMAL = 0 +NVML_FAN_FAILED = 1 _nvmlFanControlPolicy_t = c_uint NVML_FAN_POLICY_TEMPERATURE_CONTINOUS_SW = 0 -NVML_FAN_POLICY_MANUAL = 1 +NVML_FAN_POLICY_MANUAL = 1 _nvmlLedColor_t = c_uint -NVML_LED_COLOR_GREEN = 0 -NVML_LED_COLOR_AMBER = 1 +NVML_LED_COLOR_GREEN = 0 +NVML_LED_COLOR_AMBER = 1 _nvmlGpuOperationMode_t = c_uint -NVML_GOM_ALL_ON = 0 -NVML_GOM_COMPUTE = 1 -NVML_GOM_LOW_DP = 2 +NVML_GOM_ALL_ON = 0 +NVML_GOM_COMPUTE = 1 +NVML_GOM_LOW_DP = 2 _nvmlPageRetirementCause_t = c_uint NVML_PAGE_RETIREMENT_CAUSE_MULTIPLE_SINGLE_BIT_ECC_ERRORS = 0 -NVML_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR = 1 -NVML_PAGE_RETIREMENT_CAUSE_COUNT = 2 +NVML_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR = 1 +NVML_PAGE_RETIREMENT_CAUSE_COUNT = 2 _nvmlRestrictedAPI_t = c_uint -NVML_RESTRICTED_API_SET_APPLICATION_CLOCKS = 0 -NVML_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS = 1 -NVML_RESTRICTED_API_COUNT = 2 +NVML_RESTRICTED_API_SET_APPLICATION_CLOCKS = 0 +NVML_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS = 1 +NVML_RESTRICTED_API_COUNT = 2 _nvmlBridgeChipType_t = c_uint NVML_BRIDGE_CHIP_PLX = 0 @@ -366,7 +365,7 @@ NVML_TOPOLOGY_SYSTEM = 50 _nvmlGpuP2PCapsIndex_t = c_uint NVML_P2P_CAPS_INDEX_READ = 0, NVML_P2P_CAPS_INDEX_WRITE = 1 -NVML_P2P_CAPS_INDEX_NVLINK =2 +NVML_P2P_CAPS_INDEX_NVLINK = 2 NVML_P2P_CAPS_INDEX_ATOMICS = 3 # # NVML_P2P_CAPS_INDEX_PROP is deprecated. @@ -377,39 +376,39 @@ NVML_P2P_CAPS_INDEX_PCI = 4 NVML_P2P_CAPS_INDEX_UNKNOWN = 5 _nvmlGpuP2PStatus_t = c_uint -NVML_P2P_STATUS_OK = 0 +NVML_P2P_STATUS_OK = 0 NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED = 1 NVML_P2P_STATUS_CHIPSET_NOT_SUPPORTED = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED NVML_P2P_STATUS_GPU_NOT_SUPPORTED = 2 -NVML_P2P_STATUS_IOH_TOPOLOGY_NOT_SUPPORTED =3 -NVML_P2P_STATUS_DISABLED_BY_REGKEY =4 -NVML_P2P_STATUS_NOT_SUPPORTED =5 -NVML_P2P_STATUS_UNKNOWN =6 +NVML_P2P_STATUS_IOH_TOPOLOGY_NOT_SUPPORTED = 3 +NVML_P2P_STATUS_DISABLED_BY_REGKEY = 4 +NVML_P2P_STATUS_NOT_SUPPORTED = 5 +NVML_P2P_STATUS_UNKNOWN = 6 _nvmlDeviceArchitecture_t = c_uint -NVML_DEVICE_ARCH_KEPLER = 2 -NVML_DEVICE_ARCH_MAXWELL = 3 -NVML_DEVICE_ARCH_PASCAL = 4 -NVML_DEVICE_ARCH_VOLTA = 5 -NVML_DEVICE_ARCH_TURING = 6 -NVML_DEVICE_ARCH_AMPERE = 7 -NVML_DEVICE_ARCH_ADA = 8 -NVML_DEVICE_ARCH_HOPPER = 9 -NVML_DEVICE_ARCH_BLACKWELL = 10 -NVML_DEVICE_ARCH_T23X = 11 -NVML_DEVICE_ARCH_UNKNOWN = 0xffffffff +NVML_DEVICE_ARCH_KEPLER = 2 +NVML_DEVICE_ARCH_MAXWELL = 3 +NVML_DEVICE_ARCH_PASCAL = 4 +NVML_DEVICE_ARCH_VOLTA = 5 +NVML_DEVICE_ARCH_TURING = 6 +NVML_DEVICE_ARCH_AMPERE = 7 +NVML_DEVICE_ARCH_ADA = 8 +NVML_DEVICE_ARCH_HOPPER = 9 +NVML_DEVICE_ARCH_BLACKWELL = 10 +NVML_DEVICE_ARCH_T23X = 11 +NVML_DEVICE_ARCH_UNKNOWN = 0xffffffff # PCI bus Types _nvmlBusType_t = c_uint NVML_BUS_TYPE_UNKNOWN = 0 -NVML_BUS_TYPE_PCI = 1 -NVML_BUS_TYPE_PCIE = 2 -NVML_BUS_TYPE_FPCI = 3 -NVML_BUS_TYPE_AGP = 4 +NVML_BUS_TYPE_PCI = 1 +NVML_BUS_TYPE_PCIE = 2 +NVML_BUS_TYPE_FPCI = 3 +NVML_BUS_TYPE_AGP = 4 _nvmlPowerSource_t = c_uint -NVML_POWER_SOURCE_AC = 0x00000000 -NVML_POWER_SOURCE_BATTERY = 0x00000001 +NVML_POWER_SOURCE_AC = 0x00000000 +NVML_POWER_SOURCE_BATTERY = 0x00000001 NVML_POWER_SOURCE_UNDERSIZED = 0x00000002 _nvmlAdaptiveClockInfoStatus_t = c_uint @@ -418,195 +417,194 @@ NVML_ADAPTIVE_CLOCKING_INFO_STATUS_ENABLED = 0x00000001 _nvmlClockLimitId_t = c_uint NVML_CLOCK_LIMIT_ID_RANGE_START = 0xffffff00 -NVML_CLOCK_LIMIT_ID_TDP = 0xffffff01 -NVML_CLOCK_LIMIT_ID_UNLIMITED = 0xffffff02 +NVML_CLOCK_LIMIT_ID_TDP = 0xffffff01 +NVML_CLOCK_LIMIT_ID_UNLIMITED = 0xffffff02 _nvmlPcieLinkMaxSpeed_t = c_uint -NVML_PCIE_LINK_MAX_SPEED_INVALID = 0x00000000 -NVML_PCIE_LINK_MAX_SPEED_2500MBPS = 0x00000001 -NVML_PCIE_LINK_MAX_SPEED_5000MBPS = 0x00000002 -NVML_PCIE_LINK_MAX_SPEED_8000MBPS = 0x00000003 +NVML_PCIE_LINK_MAX_SPEED_INVALID = 0x00000000 +NVML_PCIE_LINK_MAX_SPEED_2500MBPS = 0x00000001 +NVML_PCIE_LINK_MAX_SPEED_5000MBPS = 0x00000002 +NVML_PCIE_LINK_MAX_SPEED_8000MBPS = 0x00000003 NVML_PCIE_LINK_MAX_SPEED_16000MBPS = 0x00000004 NVML_PCIE_LINK_MAX_SPEED_32000MBPS = 0x00000005 NVML_PCIE_LINK_MAX_SPEED_64000MBPS = 0x00000006 _nvmlPcieAtomicsCapability_t = c_uint -NVML_PCIE_ATOMICS_CAP_FETCHADD32 = 0x01 -NVML_PCIE_ATOMICS_CAP_FETCHADD64 = 0x02 -NVML_PCIE_ATOMICS_CAP_SWAP32 = 0x04 -NVML_PCIE_ATOMICS_CAP_SWAP64 = 0x08 -NVML_PCIE_ATOMICS_CAP_CAS32 = 0x10 -NVML_PCIE_ATOMICS_CAP_CAS64 = 0x20 -NVML_PCIE_ATOMICS_CAP_CAS128 = 0x40 -NVML_PCIE_ATOMICS_OPS_MAX = 7 +NVML_PCIE_ATOMICS_CAP_FETCHADD32 = 0x01 +NVML_PCIE_ATOMICS_CAP_FETCHADD64 = 0x02 +NVML_PCIE_ATOMICS_CAP_SWAP32 = 0x04 +NVML_PCIE_ATOMICS_CAP_SWAP64 = 0x08 +NVML_PCIE_ATOMICS_CAP_CAS32 = 0x10 +NVML_PCIE_ATOMICS_CAP_CAS64 = 0x20 +NVML_PCIE_ATOMICS_CAP_CAS128 = 0x40 +NVML_PCIE_ATOMICS_OPS_MAX = 7 _nvmlAffinityScope_t = c_uint -NVML_AFFINITY_SCOPE_NODE = 0 +NVML_AFFINITY_SCOPE_NODE = 0 NVML_AFFINITY_SCOPE_SOCKET = 1 _nvmlDeviceGpuRecoveryAction_t = c_uint -NVML_GPU_RECOVERY_ACTION_NONE = 0 -NVML_GPU_RECOVERY_ACTION_GPU_RESET = 1 +NVML_GPU_RECOVERY_ACTION_NONE = 0 +NVML_GPU_RECOVERY_ACTION_GPU_RESET = 1 NVML_GPU_RECOVERY_ACTION_NODE_REBOOT = 2 -NVML_GPU_RECOVERY_ACTION_DRAIN_P2P = 3 +NVML_GPU_RECOVERY_ACTION_DRAIN_P2P = 3 NVML_GPU_RECOVERY_ACTION_DRAIN_AND_RESET = 4 # C preprocessor defined values -nvmlFlagDefault = 0 -nvmlFlagForce = 1 -NVML_INIT_FLAG_NO_GPUS = 1 -NVML_INIT_FLAG_NO_ATTACH = 2 +nvmlFlagDefault = 0 +nvmlFlagForce = 1 +NVML_INIT_FLAG_NO_GPUS = 1 +NVML_INIT_FLAG_NO_ATTACH = 2 -NVML_MAX_GPC_COUNT = 32 +NVML_MAX_GPC_COUNT = 32 # buffer size -NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE = 16 -NVML_DEVICE_UUID_BUFFER_SIZE = 80 -NVML_DEVICE_UUID_V2_BUFFER_SIZE = 96 -NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE = 80 -NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE = 80 -NVML_DEVICE_NAME_BUFFER_SIZE = 64 -NVML_DEVICE_NAME_V2_BUFFER_SIZE = 96 -NVML_DEVICE_SERIAL_BUFFER_SIZE = 30 -NVML_DEVICE_PART_NUMBER_BUFFER_SIZE = 80 -NVML_DEVICE_GPU_PART_NUMBER_BUFFER_SIZE = 80 -NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE = 32 -NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE = 32 -NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE = 16 -NVML_GRID_LICENSE_BUFFER_SIZE = 128 -NVML_VGPU_NAME_BUFFER_SIZE = 64 -NVML_GRID_LICENSE_FEATURE_MAX_COUNT = 3 -NVML_VGPU_METADATA_OPAQUE_DATA_SIZE = sizeof(c_uint) + 256 -NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE = 256 -NVML_DEVICE_GPU_FRU_PART_NUMBER_BUFFER_SIZE = 0x14 # NV2080_GPU_MAX_PRODUCT_PART_NUMBER_LENGTH -NVML_PERF_MODES_BUFFER_SIZE = 2048 +NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE = 16 +NVML_DEVICE_UUID_BUFFER_SIZE = 80 +NVML_DEVICE_UUID_V2_BUFFER_SIZE = 96 +NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE = 80 +NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE = 80 +NVML_DEVICE_NAME_BUFFER_SIZE = 64 +NVML_DEVICE_NAME_V2_BUFFER_SIZE = 96 +NVML_DEVICE_SERIAL_BUFFER_SIZE = 30 +NVML_DEVICE_PART_NUMBER_BUFFER_SIZE = 80 +NVML_DEVICE_GPU_PART_NUMBER_BUFFER_SIZE = 80 +NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE = 32 +NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE = 32 +NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE = 16 +NVML_GRID_LICENSE_BUFFER_SIZE = 128 +NVML_VGPU_NAME_BUFFER_SIZE = 64 +NVML_GRID_LICENSE_FEATURE_MAX_COUNT = 3 +NVML_VGPU_METADATA_OPAQUE_DATA_SIZE = sizeof(c_uint) + 256 +NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE = 256 +NVML_DEVICE_GPU_FRU_PART_NUMBER_BUFFER_SIZE = 0x14 # NV2080_GPU_MAX_PRODUCT_PART_NUMBER_LENGTH +NVML_PERF_MODES_BUFFER_SIZE = 2048 # Format strings -NVML_DEVICE_PCI_BUS_ID_LEGACY_FMT = "%04X:%02X:%02X.0" -NVML_DEVICE_PCI_BUS_ID_FMT = "%08X:%02X:%02X.0" +NVML_DEVICE_PCI_BUS_ID_LEGACY_FMT = "%04X:%02X:%02X.0" +NVML_DEVICE_PCI_BUS_ID_FMT = "%08X:%02X:%02X.0" NVML_VALUE_NOT_AVAILABLE_ulonglong = c_ulonglong(-1) NVML_VALUE_NOT_AVAILABLE_uint = c_uint(-1) - ''' Field Identifiers. All Identifiers pertain to a device. Each ID is only used once and is guaranteed never to change. ''' -NVML_FI_DEV_ECC_CURRENT = 1 # Current ECC mode. 1=Active. 0=Inactive -NVML_FI_DEV_ECC_PENDING = 2 # Pending ECC mode. 1=Active. 0=Inactive +NVML_FI_DEV_ECC_CURRENT = 1 # Current ECC mode. 1=Active. 0=Inactive +NVML_FI_DEV_ECC_PENDING = 2 # Pending ECC mode. 1=Active. 0=Inactive #ECC Count Totals -NVML_FI_DEV_ECC_SBE_VOL_TOTAL = 3 # Total single bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_TOTAL = 4 # Total double bit volatile ECC errors -NVML_FI_DEV_ECC_SBE_AGG_TOTAL = 5 # Total single bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_TOTAL = 6 # Total double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_VOL_TOTAL = 3 # Total single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_TOTAL = 4 # Total double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_AGG_TOTAL = 5 # Total single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_TOTAL = 6 # Total double bit aggregate (persistent) ECC errors #Individual ECC locations -NVML_FI_DEV_ECC_SBE_VOL_L1 = 7 # L1 cache single bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_L1 = 8 # L1 cache double bit volatile ECC errors -NVML_FI_DEV_ECC_SBE_VOL_L2 = 9 # L2 cache single bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_L2 = 10 # L2 cache double bit volatile ECC errors -NVML_FI_DEV_ECC_SBE_VOL_DEV = 11 # Device memory single bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_DEV = 12 # Device memory double bit volatile ECC errors -NVML_FI_DEV_ECC_SBE_VOL_REG = 13 # Register file single bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_REG = 14 # Register file double bit volatile ECC errors -NVML_FI_DEV_ECC_SBE_VOL_TEX = 15 # Texture memory single bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_TEX = 16 # Texture memory double bit volatile ECC errors -NVML_FI_DEV_ECC_DBE_VOL_CBU = 17 # CBU double bit volatile ECC errors -NVML_FI_DEV_ECC_SBE_AGG_L1 = 18 # L1 cache single bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_L1 = 19 # L1 cache double bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_SBE_AGG_L2 = 20 # L2 cache single bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_L2 = 21 # L2 cache double bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_SBE_AGG_DEV = 22 # Device memory single bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_DEV = 23 # Device memory double bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_SBE_AGG_REG = 24 # Register File single bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_REG = 25 # Register File double bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_SBE_AGG_TEX = 26 # Texture memory single bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_TEX = 27 # Texture memory double bit aggregate (persistent) ECC errors -NVML_FI_DEV_ECC_DBE_AGG_CBU = 28 # CBU double bit aggregate ECC errors +NVML_FI_DEV_ECC_SBE_VOL_L1 = 7 # L1 cache single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_L1 = 8 # L1 cache double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_L2 = 9 # L2 cache single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_L2 = 10 # L2 cache double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_DEV = 11 # Device memory single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_DEV = 12 # Device memory double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_REG = 13 # Register file single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_REG = 14 # Register file double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_TEX = 15 # Texture memory single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_TEX = 16 # Texture memory double bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_CBU = 17 # CBU double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_AGG_L1 = 18 # L1 cache single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_L1 = 19 # L1 cache double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_L2 = 20 # L2 cache single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_L2 = 21 # L2 cache double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_DEV = 22 # Device memory single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_DEV = 23 # Device memory double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_REG = 24 # Register File single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_REG = 25 # Register File double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_TEX = 26 # Texture memory single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_TEX = 27 # Texture memory double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_CBU = 28 # CBU double bit aggregate ECC errors # Page Retirement -NVML_FI_DEV_RETIRED_SBE = 29 # Number of retired pages because of single bit errors -NVML_FI_DEV_RETIRED_DBE = 30 # Number of retired pages because of double bit errors -NVML_FI_DEV_RETIRED_PENDING = 31 # If any pages are pending retirement. 1=yes. 0=no. +NVML_FI_DEV_RETIRED_SBE = 29 # Number of retired pages because of single bit errors +NVML_FI_DEV_RETIRED_DBE = 30 # Number of retired pages because of double bit errors +NVML_FI_DEV_RETIRED_PENDING = 31 # If any pages are pending retirement. 1=yes. 0=no. # NvLink Flit Error Counters -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L0 = 32 # NVLink flow control CRC Error Counter for Lane 0 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L1 = 33 # NVLink flow control CRC Error Counter for Lane 1 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L2 = 34 # NVLink flow control CRC Error Counter for Lane 2 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L3 = 35 # NVLink flow control CRC Error Counter for Lane 3 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L4 = 36 # NVLink flow control CRC Error Counter for Lane 4 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L5 = 37 # NVLink flow control CRC Error Counter for Lane 5 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_TOTAL = 38 # NVLink flow control CRC Error Counter total for all Lanes +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L0 = 32 # NVLink flow control CRC Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L1 = 33 # NVLink flow control CRC Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L2 = 34 # NVLink flow control CRC Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L3 = 35 # NVLink flow control CRC Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L4 = 36 # NVLink flow control CRC Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L5 = 37 # NVLink flow control CRC Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_TOTAL = 38 # NVLink flow control CRC Error Counter total for all Lanes # NvLink CRC Data Error Counters -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L0 = 39 # NVLink data CRC Error Counter for Lane 0 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L1 = 40 # NVLink data CRC Error Counter for Lane 1 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L2 = 41 # NVLink data CRC Error Counter for Lane 2 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L3 = 42 # NVLink data CRC Error Counter for Lane 3 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L4 = 43 # NVLink data CRC Error Counter for Lane 4 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L5 = 44 # NVLink data CRC Error Counter for Lane 5 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_TOTAL = 45 # NvLink data CRC Error Counter total for all Lanes +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L0 = 39 # NVLink data CRC Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L1 = 40 # NVLink data CRC Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L2 = 41 # NVLink data CRC Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L3 = 42 # NVLink data CRC Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L4 = 43 # NVLink data CRC Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L5 = 44 # NVLink data CRC Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_TOTAL = 45 # NvLink data CRC Error Counter total for all Lanes # NvLink Replay Error Counters -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L0 = 46 # NVLink Replay Error Counter for Lane 0 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L1 = 47 # NVLink Replay Error Counter for Lane 1 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L2 = 48 # NVLink Replay Error Counter for Lane 2 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L3 = 49 # NVLink Replay Error Counter for Lane 3 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L4 = 50 # NVLink Replay Error Counter for Lane 4 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L5 = 51 # NVLink Replay Error Counter for Lane 5 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_TOTAL = 52 # NVLink Replay Error Counter total for all Lanes +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L0 = 46 # NVLink Replay Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L1 = 47 # NVLink Replay Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L2 = 48 # NVLink Replay Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L3 = 49 # NVLink Replay Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L4 = 50 # NVLink Replay Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L5 = 51 # NVLink Replay Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_TOTAL = 52 # NVLink Replay Error Counter total for all Lanes # NvLink Recovery Error Counters -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L0 = 53 # NVLink Recovery Error Counter for Lane 0 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L1 = 54 # NVLink Recovery Error Counter for Lane 1 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L2 = 55 # NVLink Recovery Error Counter for Lane 2 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L3 = 56 # NVLink Recovery Error Counter for Lane 3 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L4 = 57 # NVLink Recovery Error Counter for Lane 4 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L5 = 58 # NVLink Recovery Error Counter for Lane 5 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_TOTAL = 59 # NVLink Recovery Error Counter total for all Lanes +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L0 = 53 # NVLink Recovery Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L1 = 54 # NVLink Recovery Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L2 = 55 # NVLink Recovery Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L3 = 56 # NVLink Recovery Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L4 = 57 # NVLink Recovery Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L5 = 58 # NVLink Recovery Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_TOTAL = 59 # NVLink Recovery Error Counter total for all Lanes # NvLink Bandwidth Counters -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L0 = 60 # NVLink Bandwidth Counter for Counter Set 0, Lane 0 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L1 = 61 # NVLink Bandwidth Counter for Counter Set 0, Lane 1 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L2 = 62 # NVLink Bandwidth Counter for Counter Set 0, Lane 2 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L3 = 63 # NVLink Bandwidth Counter for Counter Set 0, Lane 3 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L4 = 64 # NVLink Bandwidth Counter for Counter Set 0, Lane 4 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L5 = 65 # NVLink Bandwidth Counter for Counter Set 0, Lane 5 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_TOTAL = 66 # NVLink Bandwidth Counter Total for Counter Set 0, All Lanes +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L0 = 60 # NVLink Bandwidth Counter for Counter Set 0, Lane 0 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L1 = 61 # NVLink Bandwidth Counter for Counter Set 0, Lane 1 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L2 = 62 # NVLink Bandwidth Counter for Counter Set 0, Lane 2 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L3 = 63 # NVLink Bandwidth Counter for Counter Set 0, Lane 3 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L4 = 64 # NVLink Bandwidth Counter for Counter Set 0, Lane 4 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L5 = 65 # NVLink Bandwidth Counter for Counter Set 0, Lane 5 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_TOTAL = 66 # NVLink Bandwidth Counter Total for Counter Set 0, All Lanes # NvLink Bandwidth Counters -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L0 = 67 # NVLink Bandwidth Counter for Counter Set 1, Lane 0 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L1 = 68 # NVLink Bandwidth Counter for Counter Set 1, Lane 1 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L2 = 69 # NVLink Bandwidth Counter for Counter Set 1, Lane 2 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L3 = 70 # NVLink Bandwidth Counter for Counter Set 1, Lane 3 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L4 = 71 # NVLink Bandwidth Counter for Counter Set 1, Lane 4 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L5 = 72 # NVLink Bandwidth Counter for Counter Set 1, Lane 5 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_TOTAL = 73 # NVLink Bandwidth Counter Total for Counter Set 1, All Lanes +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L0 = 67 # NVLink Bandwidth Counter for Counter Set 1, Lane 0 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L1 = 68 # NVLink Bandwidth Counter for Counter Set 1, Lane 1 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L2 = 69 # NVLink Bandwidth Counter for Counter Set 1, Lane 2 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L3 = 70 # NVLink Bandwidth Counter for Counter Set 1, Lane 3 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L4 = 71 # NVLink Bandwidth Counter for Counter Set 1, Lane 4 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L5 = 72 # NVLink Bandwidth Counter for Counter Set 1, Lane 5 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_TOTAL = 73 # NVLink Bandwidth Counter Total for Counter Set 1, All Lanes # Perf Policy Counters -NVML_FI_DEV_PERF_POLICY_POWER = 74 # Perf Policy Counter for Power Policy -NVML_FI_DEV_PERF_POLICY_THERMAL = 75 # Perf Policy Counter for Thermal Policy -NVML_FI_DEV_PERF_POLICY_SYNC_BOOST = 76 # Perf Policy Counter for Sync boost Policy -NVML_FI_DEV_PERF_POLICY_BOARD_LIMIT = 77 # Perf Policy Counter for Board Limit -NVML_FI_DEV_PERF_POLICY_LOW_UTILIZATION = 78 # Perf Policy Counter for Low GPU Utilization Policy -NVML_FI_DEV_PERF_POLICY_RELIABILITY = 79 # Perf Policy Counter for Reliability Policy -NVML_FI_DEV_PERF_POLICY_TOTAL_APP_CLOCKS = 80 # Perf Policy Counter for Total App Clock Policy -NVML_FI_DEV_PERF_POLICY_TOTAL_BASE_CLOCKS = 81 # Perf Policy Counter for Total Base Clocks Policy +NVML_FI_DEV_PERF_POLICY_POWER = 74 # Perf Policy Counter for Power Policy +NVML_FI_DEV_PERF_POLICY_THERMAL = 75 # Perf Policy Counter for Thermal Policy +NVML_FI_DEV_PERF_POLICY_SYNC_BOOST = 76 # Perf Policy Counter for Sync boost Policy +NVML_FI_DEV_PERF_POLICY_BOARD_LIMIT = 77 # Perf Policy Counter for Board Limit +NVML_FI_DEV_PERF_POLICY_LOW_UTILIZATION = 78 # Perf Policy Counter for Low GPU Utilization Policy +NVML_FI_DEV_PERF_POLICY_RELIABILITY = 79 # Perf Policy Counter for Reliability Policy +NVML_FI_DEV_PERF_POLICY_TOTAL_APP_CLOCKS = 80 # Perf Policy Counter for Total App Clock Policy +NVML_FI_DEV_PERF_POLICY_TOTAL_BASE_CLOCKS = 81 # Perf Policy Counter for Total Base Clocks Policy # Memory temperatures -NVML_FI_DEV_MEMORY_TEMP = 82 # Memory temperature for the device +NVML_FI_DEV_MEMORY_TEMP = 82 # Memory temperature for the device # Energy Counter -NVML_FI_DEV_TOTAL_ENERGY_CONSUMPTION = 83 # Total energy consumption for the GPU in mJ since the driver was last reloaded +NVML_FI_DEV_TOTAL_ENERGY_CONSUMPTION = 83 # Total energy consumption for the GPU in mJ since the driver was last reloaded # NVLink Speed -NVML_FI_DEV_NVLINK_SPEED_MBPS_L0 = 84 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L1 = 85 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L2 = 86 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L3 = 87 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L4 = 88 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L5 = 89 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L0 = 84 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L1 = 85 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L2 = 86 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L3 = 87 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L4 = 88 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L5 = 89 NVML_FI_DEV_NVLINK_SPEED_MBPS_COMMON = 90 # NVLink Link Count @@ -621,72 +619,72 @@ NVML_FI_DEV_PCIE_REPLAY_COUNTER = 94 NVML_FI_DEV_PCIE_REPLAY_ROLLOVER_COUNTER = 95 # NvLink Flit Error Counters -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L6 = 96 # NVLink flow control CRC Error Counter for Lane 6 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L7 = 97 # NVLink flow control CRC Error Counter for Lane 7 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L8 = 98 # NVLink flow control CRC Error Counter for Lane 8 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L9 = 99 # NVLink flow control CRC Error Counter for Lane 9 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L10 = 100 # NVLink flow control CRC Error Counter for Lane 10 -NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L11 = 101 # NVLink flow control CRC Error Counter for Lane 11 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L6 = 96 # NVLink flow control CRC Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L7 = 97 # NVLink flow control CRC Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L8 = 98 # NVLink flow control CRC Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L9 = 99 # NVLink flow control CRC Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L10 = 100 # NVLink flow control CRC Error Counter for Lane 10 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L11 = 101 # NVLink flow control CRC Error Counter for Lane 11 # NvLink CRC Data Error Counters -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L6 = 102 # NVLink data CRC Error Counter for Lane 6 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L7 = 103 # NVLink data CRC Error Counter for Lane 7 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L8 = 104 # NVLink data CRC Error Counter for Lane 8 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L9 = 105 # NVLink data CRC Error Counter for Lane 9 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L10 = 106 # NVLink data CRC Error Counter for Lane 10 -NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L11 = 107 # NVLink data CRC Error Counter for Lane 11 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L6 = 102 # NVLink data CRC Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L7 = 103 # NVLink data CRC Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L8 = 104 # NVLink data CRC Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L9 = 105 # NVLink data CRC Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L10 = 106 # NVLink data CRC Error Counter for Lane 10 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L11 = 107 # NVLink data CRC Error Counter for Lane 11 # NvLink Replay Error Counters -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L6 = 108 # NVLink Replay Error Counter for Lane 6 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L7 = 109 # NVLink Replay Error Counter for Lane 7 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L8 = 110 # NVLink Replay Error Counter for Lane 8 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L9 = 111 # NVLink Replay Error Counter for Lane 9 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L10 = 112 # NVLink Replay Error Counter for Lane 10 -NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L11 = 113 # NVLink Replay Error Counter for Lane 11 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L6 = 108 # NVLink Replay Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L7 = 109 # NVLink Replay Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L8 = 110 # NVLink Replay Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L9 = 111 # NVLink Replay Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L10 = 112 # NVLink Replay Error Counter for Lane 10 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L11 = 113 # NVLink Replay Error Counter for Lane 11 # NvLink Recovery Error Counters -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L6 = 114 # NVLink Recovery Error Counter for Lane 6 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L7 = 115 # NVLink Recovery Error Counter for Lane 7 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L8 = 116 # NVLink Recovery Error Counter for Lane 8 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L9 = 117 # NVLink Recovery Error Counter for Lane 9 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L10 = 118 # NVLink Recovery Error Counter for Lane 10 -NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L11 = 119 # NVLink Recovery Error Counter for Lane 11 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L6 = 114 # NVLink Recovery Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L7 = 115 # NVLink Recovery Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L8 = 116 # NVLink Recovery Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L9 = 117 # NVLink Recovery Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L10 = 118 # NVLink Recovery Error Counter for Lane 10 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L11 = 119 # NVLink Recovery Error Counter for Lane 11 # NvLink Bandwidth Counters -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L6 = 120 # NVLink Bandwidth Counter for Counter Set 0, Lane 6 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L7 = 121 # NVLink Bandwidth Counter for Counter Set 0, Lane 7 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L8 = 122 # NVLink Bandwidth Counter for Counter Set 0, Lane 8 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L9 = 123 # NVLink Bandwidth Counter for Counter Set 0, Lane 9 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L10 = 124 # NVLink Bandwidth Counter for Counter Set 0, Lane 10 -NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L11 = 125 # NVLink Bandwidth Counter for Counter Set 0, Lane 11 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L6 = 120 # NVLink Bandwidth Counter for Counter Set 0, Lane 6 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L7 = 121 # NVLink Bandwidth Counter for Counter Set 0, Lane 7 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L8 = 122 # NVLink Bandwidth Counter for Counter Set 0, Lane 8 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L9 = 123 # NVLink Bandwidth Counter for Counter Set 0, Lane 9 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L10 = 124 # NVLink Bandwidth Counter for Counter Set 0, Lane 10 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L11 = 125 # NVLink Bandwidth Counter for Counter Set 0, Lane 11 # NvLink Bandwidth Counters -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L6 = 126 # NVLink Bandwidth Counter for Counter Set 1, Lane 6 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L7 = 127 # NVLink Bandwidth Counter for Counter Set 1, Lane 7 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L8 = 128 # NVLink Bandwidth Counter for Counter Set 1, Lane 8 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L9 = 129 # NVLink Bandwidth Counter for Counter Set 1, Lane 9 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L10 = 130 # NVLink Bandwidth Counter for Counter Set 1, Lane 10 -NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L11 = 131 # NVLink Bandwidth Counter for Counter Set 1, Lane 11 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L6 = 126 # NVLink Bandwidth Counter for Counter Set 1, Lane 6 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L7 = 127 # NVLink Bandwidth Counter for Counter Set 1, Lane 7 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L8 = 128 # NVLink Bandwidth Counter for Counter Set 1, Lane 8 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L9 = 129 # NVLink Bandwidth Counter for Counter Set 1, Lane 9 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L10 = 130 # NVLink Bandwidth Counter for Counter Set 1, Lane 10 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L11 = 131 # NVLink Bandwidth Counter for Counter Set 1, Lane 11 # NVLink Speed -NVML_FI_DEV_NVLINK_SPEED_MBPS_L6 = 132 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L7 = 133 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L8 = 134 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L9 = 135 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L10 = 136 -NVML_FI_DEV_NVLINK_SPEED_MBPS_L11 = 137 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L6 = 132 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L7 = 133 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L8 = 134 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L9 = 135 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L10 = 136 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L11 = 137 # NVLink Throughput Counters -NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_TX = 138 # NVLink TX Data throughput in KiB -NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_RX = 139 # NVLink RX Data throughput in KiB -NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_TX = 140 # NVLink TX Data + protocol overhead in KiB -NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_RX = 141 # NVLink RX Data + protocol overhead in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_TX = 138 # NVLink TX Data throughput in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_RX = 139 # NVLink RX Data throughput in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_TX = 140 # NVLink TX Data + protocol overhead in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_RX = 141 # NVLink RX Data + protocol overhead in KiB # Row Remapper -NVML_FI_DEV_REMAPPED_COR = 142 -NVML_FI_DEV_REMAPPED_UNC = 143 -NVML_FI_DEV_REMAPPED_PENDING = 144 -NVML_FI_DEV_REMAPPED_FAILURE = 145 +NVML_FI_DEV_REMAPPED_COR = 142 +NVML_FI_DEV_REMAPPED_UNC = 143 +NVML_FI_DEV_REMAPPED_PENDING = 144 +NVML_FI_DEV_REMAPPED_FAILURE = 145 #Remote device NVLink ID NVML_FI_DEV_NVLINK_REMOTE_NVLINK_ID = 146 @@ -695,233 +693,233 @@ NVML_FI_DEV_NVLINK_REMOTE_NVLINK_ID = 146 NVML_FI_DEV_NVSWITCH_CONNECTED_LINK_COUNT = 147 # NvLink ECC Data Error Counters -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L0 = 148 #< NVLink data ECC Error Counter for Link 0 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L1 = 149 #< NVLink data ECC Error Counter for Link 1 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L2 = 150 #< NVLink data ECC Error Counter for Link 2 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L3 = 151 #< NVLink data ECC Error Counter for Link 3 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L4 = 152 #< NVLink data ECC Error Counter for Link 4 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L5 = 153 #< NVLink data ECC Error Counter for Link 5 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L6 = 154 #< NVLink data ECC Error Counter for Link 6 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L7 = 155 #< NVLink data ECC Error Counter for Link 7 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L8 = 156 #< NVLink data ECC Error Counter for Link 8 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L9 = 157 #< NVLink data ECC Error Counter for Link 9 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L10 = 158 #< NVLink data ECC Error Counter for Link 10 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L11 = 159 #< NVLink data ECC Error Counter for Link 11 -NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_TOTAL = 160 #< NvLink data ECC Error Counter total for all Links +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L0 = 148 #< NVLink data ECC Error Counter for Link 0 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L1 = 149 #< NVLink data ECC Error Counter for Link 1 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L2 = 150 #< NVLink data ECC Error Counter for Link 2 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L3 = 151 #< NVLink data ECC Error Counter for Link 3 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L4 = 152 #< NVLink data ECC Error Counter for Link 4 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L5 = 153 #< NVLink data ECC Error Counter for Link 5 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L6 = 154 #< NVLink data ECC Error Counter for Link 6 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L7 = 155 #< NVLink data ECC Error Counter for Link 7 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L8 = 156 #< NVLink data ECC Error Counter for Link 8 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L9 = 157 #< NVLink data ECC Error Counter for Link 9 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L10 = 158 #< NVLink data ECC Error Counter for Link 10 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L11 = 159 #< NVLink data ECC Error Counter for Link 11 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_TOTAL = 160 #< NvLink data ECC Error Counter total for all Links -NVML_FI_DEV_NVLINK_ERROR_DL_REPLAY = 161 -NVML_FI_DEV_NVLINK_ERROR_DL_RECOVERY = 162 -NVML_FI_DEV_NVLINK_ERROR_DL_CRC = 163 -NVML_FI_DEV_NVLINK_GET_SPEED = 164 -NVML_FI_DEV_NVLINK_GET_STATE = 165 -NVML_FI_DEV_NVLINK_GET_VERSION = 166 +NVML_FI_DEV_NVLINK_ERROR_DL_REPLAY = 161 +NVML_FI_DEV_NVLINK_ERROR_DL_RECOVERY = 162 +NVML_FI_DEV_NVLINK_ERROR_DL_CRC = 163 +NVML_FI_DEV_NVLINK_GET_SPEED = 164 +NVML_FI_DEV_NVLINK_GET_STATE = 165 +NVML_FI_DEV_NVLINK_GET_VERSION = 166 -NVML_FI_DEV_NVLINK_GET_POWER_STATE = 167 -NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD = 168 +NVML_FI_DEV_NVLINK_GET_POWER_STATE = 167 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD = 168 -NVML_FI_DEV_PCIE_L0_TO_RECOVERY_COUNTER = 169 +NVML_FI_DEV_PCIE_L0_TO_RECOVERY_COUNTER = 169 -NVML_FI_DEV_C2C_LINK_COUNT = 170 -NVML_FI_DEV_C2C_LINK_GET_STATUS = 171 -NVML_FI_DEV_C2C_LINK_GET_MAX_BW = 172 +NVML_FI_DEV_C2C_LINK_COUNT = 170 +NVML_FI_DEV_C2C_LINK_GET_STATUS = 171 +NVML_FI_DEV_C2C_LINK_GET_MAX_BW = 172 -NVML_FI_DEV_PCIE_COUNT_CORRECTABLE_ERRORS = 173 -NVML_FI_DEV_PCIE_COUNT_NAKS_RECEIVED = 174 -NVML_FI_DEV_PCIE_COUNT_RECEIVER_ERROR = 175 -NVML_FI_DEV_PCIE_COUNT_BAD_TLP = 176 -NVML_FI_DEV_PCIE_COUNT_NAKS_SENT = 177 -NVML_FI_DEV_PCIE_COUNT_BAD_DLLP = 178 -NVML_FI_DEV_PCIE_COUNT_NON_FATAL_ERROR = 179 -NVML_FI_DEV_PCIE_COUNT_FATAL_ERROR = 180 -NVML_FI_DEV_PCIE_COUNT_UNSUPPORTED_REQ = 181 -NVML_FI_DEV_PCIE_COUNT_LCRC_ERROR = 182 -NVML_FI_DEV_PCIE_COUNT_LANE_ERROR = 183 +NVML_FI_DEV_PCIE_COUNT_CORRECTABLE_ERRORS = 173 +NVML_FI_DEV_PCIE_COUNT_NAKS_RECEIVED = 174 +NVML_FI_DEV_PCIE_COUNT_RECEIVER_ERROR = 175 +NVML_FI_DEV_PCIE_COUNT_BAD_TLP = 176 +NVML_FI_DEV_PCIE_COUNT_NAKS_SENT = 177 +NVML_FI_DEV_PCIE_COUNT_BAD_DLLP = 178 +NVML_FI_DEV_PCIE_COUNT_NON_FATAL_ERROR = 179 +NVML_FI_DEV_PCIE_COUNT_FATAL_ERROR = 180 +NVML_FI_DEV_PCIE_COUNT_UNSUPPORTED_REQ = 181 +NVML_FI_DEV_PCIE_COUNT_LCRC_ERROR = 182 +NVML_FI_DEV_PCIE_COUNT_LANE_ERROR = 183 -NVML_FI_DEV_IS_RESETLESS_MIG_SUPPORTED = 184 +NVML_FI_DEV_IS_RESETLESS_MIG_SUPPORTED = 184 -NVML_FI_DEV_POWER_AVERAGE = 185 -NVML_FI_DEV_POWER_INSTANT = 186 -NVML_FI_DEV_POWER_MIN_LIMIT = 187 -NVML_FI_DEV_POWER_MAX_LIMIT = 188 -NVML_FI_DEV_POWER_DEFAULT_LIMIT = 189 -NVML_FI_DEV_POWER_CURRENT_LIMIT = 190 -NVML_FI_DEV_ENERGY = 191 -NVML_FI_DEV_POWER_REQUESTED_LIMIT = 192 +NVML_FI_DEV_POWER_AVERAGE = 185 +NVML_FI_DEV_POWER_INSTANT = 186 +NVML_FI_DEV_POWER_MIN_LIMIT = 187 +NVML_FI_DEV_POWER_MAX_LIMIT = 188 +NVML_FI_DEV_POWER_DEFAULT_LIMIT = 189 +NVML_FI_DEV_POWER_CURRENT_LIMIT = 190 +NVML_FI_DEV_ENERGY = 191 +NVML_FI_DEV_POWER_REQUESTED_LIMIT = 192 -NVML_FI_DEV_TEMPERATURE_SHUTDOWN_TLIMIT = 193 -NVML_FI_DEV_TEMPERATURE_SLOWDOWN_TLIMIT = 194 -NVML_FI_DEV_TEMPERATURE_MEM_MAX_TLIMIT = 195 -NVML_FI_DEV_TEMPERATURE_GPU_MAX_TLIMIT = 196 +NVML_FI_DEV_TEMPERATURE_SHUTDOWN_TLIMIT = 193 +NVML_FI_DEV_TEMPERATURE_SLOWDOWN_TLIMIT = 194 +NVML_FI_DEV_TEMPERATURE_MEM_MAX_TLIMIT = 195 +NVML_FI_DEV_TEMPERATURE_GPU_MAX_TLIMIT = 196 -NVML_FI_DEV_PCIE_COUNT_TX_BYTES = 197 -NVML_FI_DEV_PCIE_COUNT_RX_BYTES = 198 +NVML_FI_DEV_PCIE_COUNT_TX_BYTES = 197 +NVML_FI_DEV_PCIE_COUNT_RX_BYTES = 198 -NVML_FI_DEV_IS_MIG_MODE_INDEPENDENT_MIG_QUERY_CAPABLE = 199 +NVML_FI_DEV_IS_MIG_MODE_INDEPENDENT_MIG_QUERY_CAPABLE = 199 -NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MAX = 200 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MAX = 200 -NVML_FI_DEV_NVLINK_COUNT_XMIT_PACKETS = 201 -NVML_FI_DEV_NVLINK_COUNT_XMIT_BYTES = 202 -NVML_FI_DEV_NVLINK_COUNT_RCV_PACKETS = 203 -NVML_FI_DEV_NVLINK_COUNT_RCV_BYTES = 204 -NVML_FI_DEV_NVLINK_COUNT_VL15_DROPPED = 205 # Deprecated, do not use -NVML_FI_DEV_NVLINK_COUNT_MALFORMED_PACKET_ERRORS = 206 -NVML_FI_DEV_NVLINK_COUNT_BUFFER_OVERRUN_ERRORS = 207 -NVML_FI_DEV_NVLINK_COUNT_RCV_ERRORS = 208 -NVML_FI_DEV_NVLINK_COUNT_RCV_REMOTE_ERRORS = 209 -NVML_FI_DEV_NVLINK_COUNT_RCV_GENERAL_ERRORS = 210 -NVML_FI_DEV_NVLINK_COUNT_LOCAL_LINK_INTEGRITY_ERRORS = 211 -NVML_FI_DEV_NVLINK_COUNT_XMIT_DISCARDS = 212 +NVML_FI_DEV_NVLINK_COUNT_XMIT_PACKETS = 201 +NVML_FI_DEV_NVLINK_COUNT_XMIT_BYTES = 202 +NVML_FI_DEV_NVLINK_COUNT_RCV_PACKETS = 203 +NVML_FI_DEV_NVLINK_COUNT_RCV_BYTES = 204 +NVML_FI_DEV_NVLINK_COUNT_VL15_DROPPED = 205 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_MALFORMED_PACKET_ERRORS = 206 +NVML_FI_DEV_NVLINK_COUNT_BUFFER_OVERRUN_ERRORS = 207 +NVML_FI_DEV_NVLINK_COUNT_RCV_ERRORS = 208 +NVML_FI_DEV_NVLINK_COUNT_RCV_REMOTE_ERRORS = 209 +NVML_FI_DEV_NVLINK_COUNT_RCV_GENERAL_ERRORS = 210 +NVML_FI_DEV_NVLINK_COUNT_LOCAL_LINK_INTEGRITY_ERRORS = 211 +NVML_FI_DEV_NVLINK_COUNT_XMIT_DISCARDS = 212 NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_SUCCESSFUL_EVENTS = 213 -NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_FAILED_EVENTS = 214 -NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_EVENTS = 215 +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_FAILED_EVENTS = 214 +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_EVENTS = 215 -NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE0 = 216 # Deprecated, do not use -NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE1 = 217 # Deprecated, do not use -NVML_FI_DEV_NVLINK_COUNT_RAW_BER = 218 # Deprecated, do not use -NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_ERRORS = 219 -NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_BER = 220 -NVML_FI_DEV_NVLINK_COUNT_SYMBOL_ERRORS = 221 -NVML_FI_DEV_NVLINK_COUNT_SYMBOL_BER = 222 +NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE0 = 216 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE1 = 217 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_RAW_BER = 218 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_ERRORS = 219 +NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_BER = 220 +NVML_FI_DEV_NVLINK_COUNT_SYMBOL_ERRORS = 221 +NVML_FI_DEV_NVLINK_COUNT_SYMBOL_BER = 222 -NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MIN = 223 -NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS = 224 # Values are in the form NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_* -NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_SUPPORTED = 225 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MIN = 223 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS = 224 # Values are in the form NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_* +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_SUPPORTED = 225 -NVML_FI_DEV_RESET_STATUS = 226 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead -NVML_FI_DEV_DRAIN_AND_RESET_STATUS = 227 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead -NVML_FI_DEV_PCIE_OUTBOUND_ATOMICS_MASK = 228 -NVML_FI_DEV_PCIE_INBOUND_ATOMICS_MASK = 229 -NVML_FI_DEV_GET_GPU_RECOVERY_ACTION = 230 +NVML_FI_DEV_RESET_STATUS = 226 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead +NVML_FI_DEV_DRAIN_AND_RESET_STATUS = 227 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead +NVML_FI_DEV_PCIE_OUTBOUND_ATOMICS_MASK = 228 +NVML_FI_DEV_PCIE_INBOUND_ATOMICS_MASK = 229 +NVML_FI_DEV_GET_GPU_RECOVERY_ACTION = 230 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_0 = 235 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_1 = 236 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_2 = 237 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_3 = 238 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_4 = 239 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_5 = 240 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_6 = 241 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_7 = 242 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_8 = 243 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_9 = 244 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_10 = 245 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_11 = 246 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_12 = 247 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_13 = 248 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_14 = 249 -NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_15 = 250 -NVML_FI_PWR_SMOOTHING_ENABLED = 251 # Enablement (0/DISABLED or 1/ENABLED) -NVML_FI_PWR_SMOOTHING_PRIV_LVL = 252 # Current privilege level -NVML_FI_PWR_SMOOTHING_IMM_RAMP_DOWN_ENABLED = 253 # Immediate ramp down enablement (0/DISABLED or 1/ENABLED) -NVML_FI_PWR_SMOOTHING_APPLIED_TMP_CEIL = 254 # Applied TMP ceiling value -NVML_FI_PWR_SMOOTHING_APPLIED_TMP_FLOOR = 255 # Applied TMP floor value -NVML_FI_PWR_SMOOTHING_MAX_PERCENT_TMP_FLOOR_SETTING = 256 # Max % TMP Floor value -NVML_FI_PWR_SMOOTHING_MIN_PERCENT_TMP_FLOOR_SETTING = 257 # Min % TMP Floor value -NVML_FI_PWR_SMOOTHING_HW_CIRCUITRY_PERCENT_LIFETIME_REMAINING = 258 # HW Circuitry % lifetime remaining -NVML_FI_PWR_SMOOTHING_MAX_NUM_PRESET_PROFILES = 259 # Max number of preset profiles -NVML_FI_PWR_SMOOTHING_PROFILE_PERCENT_TMP_FLOOR = 260 # % TMP floor for a given profile -NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_UP_RATE = 261 # Ramp up rate in mW/s for a given profile -NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_RATE = 262 # Ramp down rate in mW/s for a given profile -NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_HYST_VAL = 263 # Ramp down hysteresis value in ms for a given profile -NVML_FI_PWR_SMOOTHING_ACTIVE_PRESET_PROFILE = 264 # Active preset profile number -NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_PERCENT_TMP_FLOOR = 265 # % TMP floor for a given profile -NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_UP_RATE = 266 # Ramp up rate in mW/s for a given profile -NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_RATE = 267 # Ramp down rate in mW/s for a given profile -NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_HYST_VAL = 268 # Ramp down hysteresis value in ms for a given profile +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_0 = 235 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_1 = 236 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_2 = 237 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_3 = 238 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_4 = 239 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_5 = 240 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_6 = 241 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_7 = 242 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_8 = 243 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_9 = 244 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_10 = 245 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_11 = 246 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_12 = 247 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_13 = 248 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_14 = 249 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_15 = 250 +NVML_FI_PWR_SMOOTHING_ENABLED = 251 # Enablement (0/DISABLED or 1/ENABLED) +NVML_FI_PWR_SMOOTHING_PRIV_LVL = 252 # Current privilege level +NVML_FI_PWR_SMOOTHING_IMM_RAMP_DOWN_ENABLED = 253 # Immediate ramp down enablement (0/DISABLED or 1/ENABLED) +NVML_FI_PWR_SMOOTHING_APPLIED_TMP_CEIL = 254 # Applied TMP ceiling value +NVML_FI_PWR_SMOOTHING_APPLIED_TMP_FLOOR = 255 # Applied TMP floor value +NVML_FI_PWR_SMOOTHING_MAX_PERCENT_TMP_FLOOR_SETTING = 256 # Max % TMP Floor value +NVML_FI_PWR_SMOOTHING_MIN_PERCENT_TMP_FLOOR_SETTING = 257 # Min % TMP Floor value +NVML_FI_PWR_SMOOTHING_HW_CIRCUITRY_PERCENT_LIFETIME_REMAINING = 258 # HW Circuitry % lifetime remaining +NVML_FI_PWR_SMOOTHING_MAX_NUM_PRESET_PROFILES = 259 # Max number of preset profiles +NVML_FI_PWR_SMOOTHING_PROFILE_PERCENT_TMP_FLOOR = 260 # % TMP floor for a given profile +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_UP_RATE = 261 # Ramp up rate in mW/s for a given profile +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_RATE = 262 # Ramp down rate in mW/s for a given profile +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_HYST_VAL = 263 # Ramp down hysteresis value in ms for a given profile +NVML_FI_PWR_SMOOTHING_ACTIVE_PRESET_PROFILE = 264 # Active preset profile number +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_PERCENT_TMP_FLOOR = 265 # % TMP floor for a given profile +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_UP_RATE = 266 # Ramp up rate in mW/s for a given profile +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_RATE = 267 # Ramp down rate in mW/s for a given profile +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_HYST_VAL = 268 # Ramp down hysteresis value in ms for a given profile -NVML_FI_MAX = 269 # One greater than the largest field ID defined above +NVML_FI_MAX = 269 # One greater than the largest field ID defined above # NVML_FI_DEV_NVLINK_GET_STATE state enums NVML_NVLINK_STATE_INACTIVE = 0x0 -NVML_NVLINK_STATE_ACTIVE = 0x1 -NVML_NVLINK_STATE_SLEEP = 0x2 +NVML_NVLINK_STATE_ACTIVE = 0x1 +NVML_NVLINK_STATE_SLEEP = 0x2 -NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_100US = 0 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS -NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_50US = 1 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS +NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_100US = 0 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS +NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_50US = 1 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS ## Enums needed for the method nvmlDeviceGetVirtualizationMode and nvmlDeviceSetVirtualizationMode -NVML_GPU_VIRTUALIZATION_MODE_NONE = 0 # Represents Bare Metal GPU +NVML_GPU_VIRTUALIZATION_MODE_NONE = 0 # Represents Bare Metal GPU NVML_GPU_VIRTUALIZATION_MODE_PASSTHROUGH = 1 # Device is associated with GPU-Passthorugh -NVML_GPU_VIRTUALIZATION_MODE_VGPU = 2 # Device is associated with vGPU inside virtual machine. -NVML_GPU_VIRTUALIZATION_MODE_HOST_VGPU = 3 # Device is associated with VGX hypervisor in vGPU mode -NVML_GPU_VIRTUALIZATION_MODE_HOST_VSGA = 4 # Device is associated with VGX hypervisor in vSGA mode +NVML_GPU_VIRTUALIZATION_MODE_VGPU = 2 # Device is associated with vGPU inside virtual machine. +NVML_GPU_VIRTUALIZATION_MODE_HOST_VGPU = 3 # Device is associated with VGX hypervisor in vGPU mode +NVML_GPU_VIRTUALIZATION_MODE_HOST_VSGA = 4 # Device is associated with VGX hypervisor in vSGA mode ## Lib loading ## nvmlLib = None libLoadLock = threading.Lock() -_nvmlLib_refcount = 0 # Incremented on each nvmlInit and decremented on nvmlShutdown +_nvmlLib_refcount = 0 # Incremented on each nvmlInit and decremented on nvmlShutdown ## vGPU Management -_nvmlVgpuTypeId_t = c_uint +_nvmlVgpuTypeId_t = c_uint _nvmlVgpuInstance_t = c_uint _nvmlVgpuVmIdType_t = c_uint -NVML_VGPU_VM_ID_DOMAIN_ID = 0 -NVML_VGPU_VM_ID_UUID = 1 +NVML_VGPU_VM_ID_DOMAIN_ID = 0 +NVML_VGPU_VM_ID_UUID = 1 _nvmlGridLicenseFeatureCode_t = c_uint -NVML_GRID_LICENSE_FEATURE_CODE_UNKNOWN = 0 -NVML_GRID_LICENSE_FEATURE_CODE_VGPU = 1 -NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX = 2 -NVML_GRID_LICENSE_FEATURE_CODE_VWORKSTATION = 2 # deprecated, use NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX. -NVML_GRID_LICENSE_FEATURE_CODE_GAMING = 3 -NVML_GRID_LICENSE_FEATURE_CODE_COMPUTE = 4 +NVML_GRID_LICENSE_FEATURE_CODE_UNKNOWN = 0 +NVML_GRID_LICENSE_FEATURE_CODE_VGPU = 1 +NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX = 2 +NVML_GRID_LICENSE_FEATURE_CODE_VWORKSTATION = 2 # deprecated, use NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX. +NVML_GRID_LICENSE_FEATURE_CODE_GAMING = 3 +NVML_GRID_LICENSE_FEATURE_CODE_COMPUTE = 4 _nvmlGridLicenseExpiryStatus_t = c_uint8 -NVML_GRID_LICENSE_EXPIRY_NOT_AVAILABLE = 0, # Expiry information not available -NVML_GRID_LICENSE_EXPIRY_INVALID = 1, # Invalid expiry or error fetching expiry -NVML_GRID_LICENSE_EXPIRY_VALID = 2, # Valid expiry -NVML_GRID_LICENSE_EXPIRY_NOT_APPLICABLE = 3, # Expiry not applicable -NVML_GRID_LICENSE_EXPIRY_PERMANENT = 4, # Permanent expiry +NVML_GRID_LICENSE_EXPIRY_NOT_AVAILABLE = 0, # Expiry information not available +NVML_GRID_LICENSE_EXPIRY_INVALID = 1, # Invalid expiry or error fetching expiry +NVML_GRID_LICENSE_EXPIRY_VALID = 2, # Valid expiry +NVML_GRID_LICENSE_EXPIRY_NOT_APPLICABLE = 3, # Expiry not applicable +NVML_GRID_LICENSE_EXPIRY_PERMANENT = 4, # Permanent expiry _nvmlVgpuCapability_t = c_uint -NVML_VGPU_CAP_NVLINK_P2P = 0 # vGPU P2P over NVLink is supported -NVML_VGPU_CAP_GPUDIRECT = 1 # GPUDirect capability is supported -NVML_VGPU_CAP_MULTI_VGPU_EXCLUSIVE = 2 # vGPU profile cannot be mixed with other vGPU profiles in same VM -NVML_VGPU_CAP_EXCLUSIVE_TYPE = 3 # vGPU profile cannot run on a GPU alongside other profiles of different type -NVML_VGPU_CAP_EXCLUSIVE_SIZE = 4 # vGPU profile cannot run on a GPU alongside other profiles of different size -NVML_VGPU_CAP_COUNT = 5 +NVML_VGPU_CAP_NVLINK_P2P = 0 # vGPU P2P over NVLink is supported +NVML_VGPU_CAP_GPUDIRECT = 1 # GPUDirect capability is supported +NVML_VGPU_CAP_MULTI_VGPU_EXCLUSIVE = 2 # vGPU profile cannot be mixed with other vGPU profiles in same VM +NVML_VGPU_CAP_EXCLUSIVE_TYPE = 3 # vGPU profile cannot run on a GPU alongside other profiles of different type +NVML_VGPU_CAP_EXCLUSIVE_SIZE = 4 # vGPU profile cannot run on a GPU alongside other profiles of different size +NVML_VGPU_CAP_COUNT = 5 _nvmlVgpuDriverCapability_t = c_uint -NVML_VGPU_DRIVER_CAP_HETEROGENEOUS_MULTI_VGPU = 0 # Supports mixing of different vGPU profiles within one guest VM -NVML_VGPU_DRIVER_CAP_WARM_UPDATE = 1 # Supports FSR and warm update of vGPU host driver without terminating the running guest VM -NVML_VGPU_DRIVER_CAP_COUNT = 2 +NVML_VGPU_DRIVER_CAP_HETEROGENEOUS_MULTI_VGPU = 0 # Supports mixing of different vGPU profiles within one guest VM +NVML_VGPU_DRIVER_CAP_WARM_UPDATE = 1 # Supports FSR and warm update of vGPU host driver without terminating the running guest VM +NVML_VGPU_DRIVER_CAP_COUNT = 2 _nvmlDeviceVgpuCapability_t = c_uint -NVML_DEVICE_VGPU_CAP_FRACTIONAL_MULTI_VGPU = 0 # Query whether the fractional vGPU profiles on this GPU can be used in multi-vGPU configurations -NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_PROFILES = 1 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing types -NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_SIZES = 2 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing framebuffer sizes -NVML_DEVICE_VGPU_CAP_READ_DEVICE_BUFFER_BW = 3 # Query the GPU's read_device_buffer expected bandwidth capacity in megabytes per second -NVML_DEVICE_VGPU_CAP_WRITE_DEVICE_BUFFER_BW = 4 # Query the GPU's write_device_buffer expected bandwidth capacity in megabytes per second -NVML_DEVICE_VGPU_CAP_DEVICE_STREAMING = 5 # Query whether the vGPU profiles on the GPU supports migration data streaming -NVML_DEVICE_VGPU_CAP_MINI_QUARTER_GPU = 6 # Set/Get support of mini-quarter vGPU profiles -NVML_DEVICE_VGPU_CAP_COMPUTE_MEDIA_ENGINE_GPU = 7 # Set/Get support for compute media engine vGPU profiles -NVML_DEVICE_VGPU_CAP_WARM_UPDATE = 8 # Query whether the GPU supports FSR and warm update -NVML_DEVICE_VGPU_CAP_HOMOGENEOUS_PLACEMENTS = 9 # Query whether the GPU supports reporting of placements of timesliced vGPU profiles with identical framebuffer sizes -NVML_DEVICE_VGPU_CAP_COUNT = 10 +NVML_DEVICE_VGPU_CAP_FRACTIONAL_MULTI_VGPU = 0 # Query whether the fractional vGPU profiles on this GPU can be used in multi-vGPU configurations +NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_PROFILES = 1 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing types +NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_SIZES = 2 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing framebuffer sizes +NVML_DEVICE_VGPU_CAP_READ_DEVICE_BUFFER_BW = 3 # Query the GPU's read_device_buffer expected bandwidth capacity in megabytes per second +NVML_DEVICE_VGPU_CAP_WRITE_DEVICE_BUFFER_BW = 4 # Query the GPU's write_device_buffer expected bandwidth capacity in megabytes per second +NVML_DEVICE_VGPU_CAP_DEVICE_STREAMING = 5 # Query whether the vGPU profiles on the GPU supports migration data streaming +NVML_DEVICE_VGPU_CAP_MINI_QUARTER_GPU = 6 # Set/Get support of mini-quarter vGPU profiles +NVML_DEVICE_VGPU_CAP_COMPUTE_MEDIA_ENGINE_GPU = 7 # Set/Get support for compute media engine vGPU profiles +NVML_DEVICE_VGPU_CAP_WARM_UPDATE = 8 # Query whether the GPU supports FSR and warm update +NVML_DEVICE_VGPU_CAP_HOMOGENEOUS_PLACEMENTS = 9 # Query whether the GPU supports reporting of placements of timesliced vGPU profiles with identical framebuffer sizes +NVML_DEVICE_VGPU_CAP_COUNT = 10 _nvmlVgpuGuestInfoState_t = c_uint NVML_VGPU_INSTANCE_GUEST_INFO_STATE_UNINITIALIZED = 0 -NVML_VGPU_INSTANCE_GUEST_INFO_STATE_INITIALIZED = 1 +NVML_VGPU_INSTANCE_GUEST_INFO_STATE_INITIALIZED = 1 _nvmlVgpuVmCompatibility_t = c_uint -NVML_VGPU_VM_COMPATIBILITY_NONE = 0x0 -NVML_VGPU_VM_COMPATIBILITY_COLD = 0x1 -NVML_VGPU_VM_COMPATIBILITY_HIBERNATE = 0x2 -NVML_VGPU_VM_COMPATIBILITY_SLEEP = 0x4 -NVML_VGPU_VM_COMPATIBILITY_LIVE = 0x8 +NVML_VGPU_VM_COMPATIBILITY_NONE = 0x0 +NVML_VGPU_VM_COMPATIBILITY_COLD = 0x1 +NVML_VGPU_VM_COMPATIBILITY_HIBERNATE = 0x2 +NVML_VGPU_VM_COMPATIBILITY_SLEEP = 0x4 +NVML_VGPU_VM_COMPATIBILITY_LIVE = 0x8 _nvmlVgpuPgpuCompatibilityLimitCode_t = c_uint -NVML_VGPU_COMPATIBILITY_LIMIT_NONE = 0x0 -NVML_VGPU_COMPATIBILITY_LIMIT_HOST_DRIVER = 0x1 -NVML_VGPU_COMPATIBILITY_LIMIT_GUEST_DRIVER = 0x2 -NVML_VGPU_COMPATIBILITY_LIMIT_GPU = 0x4 -NVML_VGPU_COMPATIBILITY_LIMIT_OTHER = 0x80000000 +NVML_VGPU_COMPATIBILITY_LIMIT_NONE = 0x0 +NVML_VGPU_COMPATIBILITY_LIMIT_HOST_DRIVER = 0x1 +NVML_VGPU_COMPATIBILITY_LIMIT_GUEST_DRIVER = 0x2 +NVML_VGPU_COMPATIBILITY_LIMIT_GPU = 0x4 +NVML_VGPU_COMPATIBILITY_LIMIT_OTHER = 0x80000000 _nvmlHostVgpuMode_t = c_uint -NVML_HOST_VGPU_MODE_NON_SRIOV = 0 -NVML_HOST_VGPU_MODE_SRIOV = 1 +NVML_HOST_VGPU_MODE_NON_SRIOV = 0 +NVML_HOST_VGPU_MODE_SRIOV = 1 _nvmlConfComputeGpusReadyState_t = c_uint NVML_CC_ACCEPTING_CLIENT_REQUESTS_FALSE = 0 @@ -944,11 +942,11 @@ NVML_CC_SYSTEM_DEVTOOLS_MODE_ON = 1 NVML_CC_SYSTEM_MULTIGPU_NONE = 0 NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE = 1 - + NVML_CC_SYSTEM_ENVIRONMENT_UNAVAILABLE = 0 NVML_CC_SYSTEM_ENVIRONMENT_SIM = 1 NVML_CC_SYSTEM_ENVIRONMENT_PROD = 2 - + _nvmlConfComputeCcFeature_t = c_uint NVML_CC_SYSTEM_FEATURE_DISABLED = 0 NVML_CC_SYSTEM_FEATURE_ENABLED = 1 @@ -960,35 +958,38 @@ NVML_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MAX = 65 # GSP firmware NVML_GSP_FIRMWARE_VERSION_BUF_SIZE = 0x40 + class NVMLLibraryMismatchError(Exception): pass + ## Error Checking ## class NVMLError(Exception): _valClassMapping = dict() # List of currently known error codes _errcode_to_string = { - NVML_ERROR_UNINITIALIZED: "Uninitialized", - NVML_ERROR_INVALID_ARGUMENT: "Invalid Argument", - NVML_ERROR_NOT_SUPPORTED: "Not Supported", - NVML_ERROR_NO_PERMISSION: "Insufficient Permissions", + NVML_ERROR_UNINITIALIZED: "Uninitialized", + NVML_ERROR_INVALID_ARGUMENT: "Invalid Argument", + NVML_ERROR_NOT_SUPPORTED: "Not Supported", + NVML_ERROR_NO_PERMISSION: "Insufficient Permissions", NVML_ERROR_ALREADY_INITIALIZED: "Already Initialized", - NVML_ERROR_NOT_FOUND: "Not Found", - NVML_ERROR_INSUFFICIENT_SIZE: "Insufficient Size", - NVML_ERROR_INSUFFICIENT_POWER: "Insufficient External Power", - NVML_ERROR_DRIVER_NOT_LOADED: "Driver Not Loaded", - NVML_ERROR_TIMEOUT: "Timeout", - NVML_ERROR_IRQ_ISSUE: "Interrupt Request Issue", - NVML_ERROR_LIBRARY_NOT_FOUND: "NVML Shared Library Not Found", - NVML_ERROR_FUNCTION_NOT_FOUND: "Function Not Found", - NVML_ERROR_CORRUPTED_INFOROM: "Corrupted infoROM", - NVML_ERROR_GPU_IS_LOST: "GPU is lost", - NVML_ERROR_RESET_REQUIRED: "GPU requires restart", - NVML_ERROR_OPERATING_SYSTEM: "The operating system has blocked the request.", + NVML_ERROR_NOT_FOUND: "Not Found", + NVML_ERROR_INSUFFICIENT_SIZE: "Insufficient Size", + NVML_ERROR_INSUFFICIENT_POWER: "Insufficient External Power", + NVML_ERROR_DRIVER_NOT_LOADED: "Driver Not Loaded", + NVML_ERROR_TIMEOUT: "Timeout", + NVML_ERROR_IRQ_ISSUE: "Interrupt Request Issue", + NVML_ERROR_LIBRARY_NOT_FOUND: "NVML Shared Library Not Found", + NVML_ERROR_FUNCTION_NOT_FOUND: "Function Not Found", + NVML_ERROR_CORRUPTED_INFOROM: "Corrupted infoROM", + NVML_ERROR_GPU_IS_LOST: "GPU is lost", + NVML_ERROR_RESET_REQUIRED: "GPU requires restart", + NVML_ERROR_OPERATING_SYSTEM: "The operating system has blocked the request.", NVML_ERROR_LIB_RM_VERSION_MISMATCH: "RM has detected an NVML/RM version mismatch.", - NVML_ERROR_MEMORY: "Insufficient Memory", - NVML_ERROR_UNKNOWN: "Unknown Error", - } + NVML_ERROR_MEMORY: "Insufficient Memory", + NVML_ERROR_UNKNOWN: "Unknown Error", + } + def __new__(typ, value): ''' Maps value to a proper subclass of NVMLError. @@ -999,6 +1000,7 @@ class NVMLError(Exception): obj = Exception.__new__(typ) obj.value = value return obj + def __str__(self): try: if self.value not in NVMLError._errcode_to_string: @@ -1006,14 +1008,17 @@ class NVMLError(Exception): return NVMLError._errcode_to_string[self.value] except NVMLError: return "NVML Error with code %d" % self.value + def __eq__(self, other): return self.value == other.value + def nvmlExceptionClass(nvmlErrorCode): if nvmlErrorCode not in NVMLError._valClassMapping: raise ValueError('nvmlErrorCode %s is not valid' % nvmlErrorCode) return NVMLError._valClassMapping[nvmlErrorCode] + def _extractNVMLErrorsAsClasses(): ''' Generates a hierarchy of classes on top of NVMLError class. @@ -1030,24 +1035,34 @@ def _extractNVMLErrorsAsClasses(): # e.g. Turn NVML_ERROR_ALREADY_INITIALIZED into NVMLError_AlreadyInitialized class_name = "NVMLError_" + string.capwords(err_name.replace("NVML_ERROR_", ""), "_").replace("_", "") err_val = getattr(this_module, err_name) + def gen_new(val): + def new(typ): obj = NVMLError.__new__(typ, val) return obj + return new - new_error_class = type(class_name, (NVMLError,), {'__new__': gen_new(err_val)}) + + new_error_class = type(class_name, (NVMLError, ), {'__new__': gen_new(err_val)}) new_error_class.__module__ = __name__ setattr(this_module, class_name, new_error_class) NVMLError._valClassMapping[err_val] = new_error_class + + _extractNVMLErrorsAsClasses() + def _nvmlCheckReturn(ret): if (ret != NVML_SUCCESS): raise NVMLError(ret) return ret + ## Function access ## -_nvmlGetFunctionPointer_cache = dict() # function pointers are cached to prevent unnecessary libLoadLock locking +_nvmlGetFunctionPointer_cache = dict() # function pointers are cached to prevent unnecessary libLoadLock locking + + def _nvmlGetFunctionPointer(name): global nvmlLib @@ -1068,17 +1083,21 @@ def _nvmlGetFunctionPointer(name): # lock is always freed libLoadLock.release() + ## Alternative object # Allows the object to be printed # Allows mismatched types to be assigned # - like None when the Structure variant requires c_uint class nvmlFriendlyObject(object): + def __init__(self, dictionary): for x in dictionary: setattr(self, x, dictionary[x]) + def __str__(self): return self.__dict__.__str__() + def nvmlStructToFriendlyObject(struct): d = {} for x in struct._fields_: @@ -1089,23 +1108,28 @@ def nvmlStructToFriendlyObject(struct): obj = nvmlFriendlyObject(d) return obj + # pack the object so it can be passed to the NVML library def nvmlFriendlyObjectToStruct(obj, model): for x in model._fields_: key = x[0] value = obj.__dict__[key] # any c_char_p in python3 needs to be bytes, default encoding works fine. - if sys.version_info >= (3,): + if sys.version_info >= (3, ): setattr(model, key, value.encode()) else: setattr(model, key, value) return model + ## Unit structures class struct_c_nvmlUnit_t(Structure): - pass # opaque handle + pass # opaque handle + + c_nvmlUnit_t = POINTER(struct_c_nvmlUnit_t) + class _PrintableStructure(Structure): """ Abstract class that produces nicer __str__ output than ctypes.Structure. @@ -1126,6 +1150,7 @@ class _PrintableStructure(Structure): Exact format of returned str from this class is subject to change in the future. """ _fmt_ = {} + def __str__(self): result = [] for x in self._fields_: @@ -1137,7 +1162,7 @@ class _PrintableStructure(Structure): elif "" in self._fmt_: fmt = self._fmt_[""] result.append(("%s: " + fmt) % (key, value)) - return self.__class__.__name__ + "(" + ", ".join(result) + ")" + return self.__class__.__name__ + "(" + ", ".join(result) + ")" def __getattribute__(self, name): res = super(_PrintableStructure, self).__getattribute__(name) @@ -1158,6 +1183,7 @@ class _PrintableStructure(Structure): value = value.encode() super(_PrintableStructure, self).__setattr__(name, value) + class c_nvmlUnitInfo_t(_PrintableStructure): _fields_ = [ ('name', c_char * 96), @@ -1166,12 +1192,13 @@ class c_nvmlUnitInfo_t(_PrintableStructure): ('firmwareVersion', c_char * 96), ] -class c_nvmlC2cModeInfo_v1_t(_PrintableStructure): - _fields_ = [ - ('isC2cEnabled', c_uint) - ] -nvmlC2cModeInfo_v1 = 0x1000008; +class c_nvmlC2cModeInfo_v1_t(_PrintableStructure): + _fields_ = [('isC2cEnabled', c_uint)] + + +nvmlC2cModeInfo_v1 = 0x1000008 + class c_nvmlLedState_t(_PrintableStructure): _fields_ = [ @@ -1179,6 +1206,7 @@ class c_nvmlLedState_t(_PrintableStructure): ('color', _nvmlLedColor_t), ] + class c_nvmlPSUInfo_t(_PrintableStructure): _fields_ = [ ('state', c_char * 256), @@ -1187,23 +1215,26 @@ class c_nvmlPSUInfo_t(_PrintableStructure): ('power', c_uint), ] + class c_nvmlUnitFanInfo_t(_PrintableStructure): _fields_ = [ ('speed', c_uint), ('state', _nvmlFanState_t), ] + class c_nvmlUnitFanSpeeds_t(_PrintableStructure): - _fields_ = [ - ('fans', c_nvmlUnitFanInfo_t * 24), - ('count', c_uint) - ] + _fields_ = [('fans', c_nvmlUnitFanInfo_t * 24), ('count', c_uint)] + ## Device structures class struct_c_nvmlDevice_t(Structure): - pass # opaque handle + pass # opaque handle + + c_nvmlDevice_t = POINTER(struct_c_nvmlDevice_t) + class nvmlPciInfoExt_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -1217,18 +1248,20 @@ class nvmlPciInfoExt_v1_t(_PrintableStructure): ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), ] _fmt_ = { - 'version' : "0x%04X", - 'domain' : "0x%04X", - 'bus' : "0x%02X", - 'device' : "0x%02X", - 'pciDeviceId' : "0x%08X", - 'pciSubSystemId' : "0x%08X", - 'baseClass' : "0x%01X", - 'subClass' : "0x%01X", - } + 'version': "0x%04X", + 'domain': "0x%04X", + 'bus': "0x%02X", + 'device': "0x%02X", + 'pciDeviceId': "0x%08X", + 'pciSubSystemId': "0x%08X", + 'baseClass': "0x%01X", + 'subClass': "0x%01X", + } + nvmlPciInfoExt_v1 = 0x1000040 + # Legacy pciInfo used for _v1 and _v2 class nvmlPciInfo_v2_t(_PrintableStructure): _fields_ = [ @@ -1246,12 +1279,13 @@ class nvmlPciInfo_v2_t(_PrintableStructure): ('reserved3', c_uint), ] _fmt_ = { - 'domain' : "0x%04X", - 'bus' : "0x%02X", - 'device' : "0x%02X", - 'pciDeviceId' : "0x%08X", - 'pciSubSystemId' : "0x%08X", - } + 'domain': "0x%04X", + 'bus': "0x%02X", + 'device': "0x%02X", + 'pciDeviceId': "0x%08X", + 'pciSubSystemId': "0x%08X", + } + class nvmlPciInfo_t(_PrintableStructure): _fields_ = [ @@ -1269,12 +1303,13 @@ class nvmlPciInfo_t(_PrintableStructure): ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), ] _fmt_ = { - 'domain' : "0x%08X", - 'bus' : "0x%02X", - 'device' : "0x%02X", - 'pciDeviceId' : "0x%08X", - 'pciSubSystemId' : "0x%08X", - } + 'domain': "0x%08X", + 'bus': "0x%02X", + 'device': "0x%02X", + 'pciDeviceId': "0x%08X", + 'pciSubSystemId': "0x%08X", + } + class c_nvmlSystemDriverBranchInfo_v1_t(_PrintableStructure): _fields_ = [ @@ -1282,13 +1317,13 @@ class c_nvmlSystemDriverBranchInfo_v1_t(_PrintableStructure): ("branch", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), ] + SystemDriverBranchInfo_v1 = 0x1000054 + class c_nvmlExcludedDeviceInfo_t(_PrintableStructure): - _fields_ = [ - ('pci', nvmlPciInfo_t), - ('uuid', c_char * NVML_DEVICE_UUID_BUFFER_SIZE) - ] + _fields_ = [('pci', nvmlPciInfo_t), ('uuid', c_char * NVML_DEVICE_UUID_BUFFER_SIZE)] + class nvmlNvLinkUtilizationControl_t(_PrintableStructure): _fields_ = [ @@ -1296,6 +1331,7 @@ class nvmlNvLinkUtilizationControl_t(_PrintableStructure): ('pktfilter', _nvmlNvLinkUtilizationCountPktTypes_t), ] + class c_nvmlMemory_t(_PrintableStructure): _fields_ = [ ('total', c_ulonglong), @@ -1304,6 +1340,7 @@ class c_nvmlMemory_t(_PrintableStructure): ] _fmt_ = {'': "%d B"} + class c_nvmlMemory_v2_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -1314,8 +1351,10 @@ class c_nvmlMemory_v2_t(_PrintableStructure): ] _fmt_ = {'': "%d B"} + nvmlMemory_v2 = 0x02000028 + class c_nvmlBAR1Memory_t(_PrintableStructure): _fields_ = [ ('bar1Total', c_ulonglong), @@ -1324,18 +1363,18 @@ class c_nvmlBAR1Memory_t(_PrintableStructure): ] _fmt_ = {'': "%d B"} + class nvmlClkMonFaultInfo_t(Structure): - _fields_ = [("clkApiDomain", c_uint), - ("clkDomainFaultMask", c_uint) - ] + _fields_ = [("clkApiDomain", c_uint), ("clkDomainFaultMask", c_uint)] + MAX_CLK_DOMAINS = 32 + class nvmlClkMonStatus_t(Structure): - _fields_ = [("bGlobalStatus", c_uint), - ("clkMonListSize", c_uint), - ("clkMonList", nvmlClkMonFaultInfo_t * MAX_CLK_DOMAINS) - ] + _fields_ = [("bGlobalStatus", c_uint), ("clkMonListSize", c_uint), + ("clkMonList", nvmlClkMonFaultInfo_t * MAX_CLK_DOMAINS)] + # On Windows with the WDDM driver, usedGpuMemory is reported as None # Code that processes this structure should check for None, I.E. @@ -1357,14 +1396,16 @@ class c_nvmlProcessInfo_v2_t(_PrintableStructure): ] _fmt_ = {'usedGpuMemory': "%d B"} + c_nvmlProcessInfo_v3_t = c_nvmlProcessInfo_v2_t c_nvmlProcessInfo_t = c_nvmlProcessInfo_v3_t _nvmlProcessMode_t = c_uint -NVML_PROCESS_MODE_COMPUTE = 0 +NVML_PROCESS_MODE_COMPUTE = 0 NVML_PROCESS_MODE_GRAPHICS = 1 -NVML_PROCESS_MODE_MPS = 2 +NVML_PROCESS_MODE_MPS = 2 + class c_nvmlProcessDetail_v1_t(Structure): _fields_ = [ @@ -1375,6 +1416,7 @@ class c_nvmlProcessDetail_v1_t(Structure): ('usedGpuCcProtectedMemory', c_ulonglong), ] + class c_nvmlProcessDetailList_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -1384,22 +1426,26 @@ class c_nvmlProcessDetailList_v1_t(_PrintableStructure): ] _fmt_ = {'numProcArrayEntries': "%d B"} + c_nvmlProcessDetailList_t = c_nvmlProcessDetailList_v1_t nvmlProcessDetailList_v1 = 0x1000018 + class c_nvmlBridgeChipInfo_t(_PrintableStructure): _fields_ = [ ('type', _nvmlBridgeChipType_t), ('fwVersion', c_uint), ] + class c_nvmlBridgeChipHierarchy_t(_PrintableStructure): _fields_ = [ ('bridgeCount', c_uint), ('bridgeChipInfo', c_nvmlBridgeChipInfo_t * 128), ] + class c_nvmlEccErrorCounts_t(_PrintableStructure): _fields_ = [ ('l1Cache', c_ulonglong), @@ -1408,6 +1454,7 @@ class c_nvmlEccErrorCounts_t(_PrintableStructure): ('registerFile', c_ulonglong), ] + class c_nvmlUtilization_t(_PrintableStructure): _fields_ = [ ('gpu', c_uint), @@ -1415,6 +1462,7 @@ class c_nvmlUtilization_t(_PrintableStructure): ] _fmt_ = {'': "%d %%"} + # Added in 2.285 class c_nvmlHwbcEntry_t(_PrintableStructure): _fields_ = [ @@ -1422,6 +1470,7 @@ class c_nvmlHwbcEntry_t(_PrintableStructure): ('firmwareVersion', c_char * 32), ] + class c_nvmlValue_t(Union): _fields_ = [ ('dVal', c_double), @@ -1433,80 +1482,79 @@ class c_nvmlValue_t(Union): ('usVal', c_ushort), ] + class c_nvmlSample_t(_PrintableStructure): _fields_ = [ ('timeStamp', c_ulonglong), ('sampleValue', c_nvmlValue_t), ] + class c_nvmlViolationTime_t(_PrintableStructure): _fields_ = [ ('referenceTime', c_ulonglong), ('violationTime', c_ulonglong), ] + class c_nvmlFieldValue_t(_PrintableStructure): - _fields_ = [ - ('fieldId', c_uint32), - ('scopeId', c_uint32), - ('timestamp', c_int64), - ('latencyUsec', c_int64), - ('valueType', _nvmlValueType_t), - ('nvmlReturn', _nvmlReturn_t), - ('value', c_nvmlValue_t) - ] + _fields_ = [('fieldId', c_uint32), ('scopeId', c_uint32), ('timestamp', c_int64), ('latencyUsec', c_int64), + ('valueType', _nvmlValueType_t), ('nvmlReturn', _nvmlReturn_t), ('value', c_nvmlValue_t)] + NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES = 23 nvmlNvlinkSupportedBwModes_v1 = 0x100001c + + class c_nvmlNvlinkSupportedBwModes_v1_t(_PrintableStructure): - _fields_ = [ - ('version', c_uint), - ('bwModes', c_uint8 * NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES), - ('totalBwModes', c_uint8) - ] + _fields_ = [('version', c_uint), ('bwModes', c_uint8 * NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES), + ('totalBwModes', c_uint8)] def __init__(self): super(c_nvmlNvlinkSupportedBwModes_v1_t, self).__init__(version=nvmlNvlinkSupportedBwModes_v1) + nvmlNvlinkGetBwMode_v1 = 0x100000c + + class c_nvmlNvlinkGetBwMode_v1_t(_PrintableStructure): - _fields_ = [ - ('version', c_uint), - ('bIsBest', c_uint), - ('bwMode', c_uint8) - ] + _fields_ = [('version', c_uint), ('bIsBest', c_uint), ('bwMode', c_uint8)] def __init__(self): super(c_nvmlNvlinkGetBwMode_v1_t, self).__init__(version=nvmlNvlinkGetBwMode_v1) + nvmlNvlinkSetBwMode_v1 = 0x100000c + + class c_nvmlNvlinkSetBwMode_v1_t(_PrintableStructure): - _fields_ = [ - ('version', c_uint), - ('bSetBest', c_uint), - ('bwMode', c_uint8) - ] + _fields_ = [('version', c_uint), ('bSetBest', c_uint), ('bwMode', c_uint8)] def __init__(self): super(c_nvmlNvlinkSetBwMode_v1_t, self).__init__(version=nvmlNvlinkSetBwMode_v1) + class c_nvmlVgpuHeterogeneousMode_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), ('mode', c_uint), ] + VgpuHeterogeneousMode_v1 = 0x1000008 + class c_nvmlVgpuPlacementId_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), ('placementId', c_uint), ] + VgpuPlacementId_v1 = 0x1000008 + class c_nvmlVgpuPlacementList_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -1515,10 +1563,12 @@ class c_nvmlVgpuPlacementList_v1_t(_PrintableStructure): ('placementIds', POINTER(c_uint)), ] + VgpuPlacementList_v1 = 0x1000018 -NVML_VGPU_PGPU_HETEROGENEOUS_MODE = 0 -NVML_VGPU_PGPU_HOMOGENEOUS_MODE = 1 +NVML_VGPU_PGPU_HETEROGENEOUS_MODE = 0 +NVML_VGPU_PGPU_HOMOGENEOUS_MODE = 1 + class c_nvmlVgpuPlacementList_v2_t(_PrintableStructure): _fields_ = [ @@ -1529,16 +1579,20 @@ class c_nvmlVgpuPlacementList_v2_t(_PrintableStructure): ('mode', c_uint), ] + VgpuPlacementList_v2 = 0x2000020 + class c_nvmlVgpuTypeBar1Info_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), ('bar1Size', c_ulonglong), ] + VgpuTypeBar1Info_v1 = 0x1000010 + class c_nvmlVgpuInstanceUtilizationSample_t(_PrintableStructure): _fields_ = [ ('vgpuInstance', _nvmlVgpuInstance_t), @@ -1549,6 +1603,7 @@ class c_nvmlVgpuInstanceUtilizationSample_t(_PrintableStructure): ('decUtil', c_nvmlValue_t), ] + class c_nvmlVgpuInstanceUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ ('timeStamp', c_ulonglong), @@ -1561,6 +1616,7 @@ class c_nvmlVgpuInstanceUtilizationInfo_v1_t(_PrintableStructure): ('ofaUtil', c_nvmlValue_t), ] + class c_nvmlVgpuInstancesUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -1570,8 +1626,10 @@ class c_nvmlVgpuInstancesUtilizationInfo_v1_t(_PrintableStructure): ('vgpuUtilArray', POINTER(c_nvmlVgpuInstanceUtilizationInfo_v1_t)), ] + VgpuInstancesUtilizationInfo_v1 = 0x01000020 + class c_nvmlVgpuProcessUtilizationSample_t(_PrintableStructure): _fields_ = [ ('vgpuInstance', _nvmlVgpuInstance_t), @@ -1584,6 +1642,7 @@ class c_nvmlVgpuProcessUtilizationSample_t(_PrintableStructure): ('decUtil', c_uint), ] + class c_nvmlVgpuProcessUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ ('processName', c_char * NVML_VGPU_NAME_BUFFER_SIZE), @@ -1598,6 +1657,7 @@ class c_nvmlVgpuProcessUtilizationInfo_v1_t(_PrintableStructure): ('ofaUtil', c_uint), ] + class c_nvmlVgpuProcessesUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -1606,41 +1666,48 @@ class c_nvmlVgpuProcessesUtilizationInfo_v1_t(_PrintableStructure): ('vgpuProcUtilArray', POINTER(c_nvmlVgpuProcessUtilizationInfo_v1_t)), ] + VgpuProcessesUtilizationInfo_v1 = 0x01000018 + class nvmlVgpuRuntimeState_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), ('size', c_ulonglong), ] + VgpuRuntimeState_v1 = 0x1000010 + class c_nvmlVgpuLicenseExpiry_t(_PrintableStructure): _fields_ = [ - ('year', c_uint32), - ('month', c_uint16), - ('day', c_uint16), - ('hour', c_uint16), - ('min', c_uint16), - ('sec', c_uint16), - ('status', c_uint8), + ('year', c_uint32), + ('month', c_uint16), + ('day', c_uint16), + ('hour', c_uint16), + ('min', c_uint16), + ('sec', c_uint16), + ('status', c_uint8), ] -NVML_GRID_LICENSE_STATE_UNKNOWN = 0 -NVML_GRID_LICENSE_STATE_UNINITIALIZED = 1 + +NVML_GRID_LICENSE_STATE_UNKNOWN = 0 +NVML_GRID_LICENSE_STATE_UNINITIALIZED = 1 NVML_GRID_LICENSE_STATE_UNLICENSED_UNRESTRICTED = 2 -NVML_GRID_LICENSE_STATE_UNLICENSED_RESTRICTED = 3 -NVML_GRID_LICENSE_STATE_UNLICENSED = 4 -NVML_GRID_LICENSE_STATE_LICENSED = 5 +NVML_GRID_LICENSE_STATE_UNLICENSED_RESTRICTED = 3 +NVML_GRID_LICENSE_STATE_UNLICENSED = 4 +NVML_GRID_LICENSE_STATE_LICENSED = 5 + class c_nvmlVgpuLicenseInfo_t(_PrintableStructure): _fields_ = [ - ('isLicensed', c_uint8), - ('licenseExpiry', c_nvmlVgpuLicenseExpiry_t), - ('currentState', c_uint), + ('isLicensed', c_uint8), + ('licenseExpiry', c_nvmlVgpuLicenseExpiry_t), + ('currentState', c_uint), ] + class c_nvmlEncoderSession_t(_PrintableStructure): _fields_ = [ ('sessionId', c_uint), @@ -1653,6 +1720,7 @@ class c_nvmlEncoderSession_t(_PrintableStructure): ('encodeLatency', c_uint), ] + class c_nvmlProcessUtilizationSample_t(_PrintableStructure): _fields_ = [ ('pid', c_uint), @@ -1663,6 +1731,7 @@ class c_nvmlProcessUtilizationSample_t(_PrintableStructure): ('decUtil', c_uint), ] + class c_nvmlProcessUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ ('timeStamp', c_ulonglong), @@ -1675,6 +1744,7 @@ class c_nvmlProcessUtilizationInfo_v1_t(_PrintableStructure): ('ofaUtil', c_uint), ] + class c_nvmlProcessesUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -1683,36 +1753,41 @@ class c_nvmlProcessesUtilizationInfo_v1_t(_PrintableStructure): ('procUtilArray', POINTER(c_nvmlProcessUtilizationInfo_v1_t)), ] + ProcessesUtilizationInfo_v1 = 0x01000018 + class c_nvmlGridLicenseExpiry_t(_PrintableStructure): _fields_ = [ - ('year', c_uint32), - ('month', c_uint16), - ('day', c_uint16), - ('hour', c_uint16), - ('min', c_uint16), - ('sec', c_uint16), - ('status', c_uint8), + ('year', c_uint32), + ('month', c_uint16), + ('day', c_uint16), + ('hour', c_uint16), + ('min', c_uint16), + ('sec', c_uint16), + ('status', c_uint8), ] + class c_nvmlGridLicensableFeature_v4_t(_PrintableStructure): _fields_ = [ - ('featureCode', _nvmlGridLicenseFeatureCode_t), - ('featureState', c_uint), - ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), - ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ('featureCode', _nvmlGridLicenseFeatureCode_t), + ('featureState', c_uint), + ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), ('featureEnabled', c_uint), - ('licenseExpiry', c_nvmlGridLicenseExpiry_t), + ('licenseExpiry', c_nvmlGridLicenseExpiry_t), ] + class c_nvmlGridLicensableFeatures_v4_t(_PrintableStructure): _fields_ = [ - ('isGridLicenseSupported', c_int), + ('isGridLicenseSupported', c_int), ('licensableFeaturesCount', c_uint), - ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v4_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v4_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), ] + class c_nvmlGridLicensableFeature_v3_t(_PrintableStructure): _fields_ = [ ('featureCode', _nvmlGridLicenseFeatureCode_t), @@ -1722,6 +1797,7 @@ class c_nvmlGridLicensableFeature_v3_t(_PrintableStructure): ('featureEnabled', c_uint), ] + class c_nvmlGridLicensableFeatures_v3_t(_PrintableStructure): _fields_ = [ ('isGridLicenseSupported', c_int), @@ -1729,6 +1805,7 @@ class c_nvmlGridLicensableFeatures_v3_t(_PrintableStructure): ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v3_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), ] + class c_nvmlGridLicensableFeature_v2_t(_PrintableStructure): _fields_ = [ ('featureCode', _nvmlGridLicenseFeatureCode_t), @@ -1737,6 +1814,7 @@ class c_nvmlGridLicensableFeature_v2_t(_PrintableStructure): ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), ] + class c_nvmlGridLicensableFeatures_v2_t(_PrintableStructure): _fields_ = [ ('isGridLicenseSupported', c_int), @@ -1744,6 +1822,7 @@ class c_nvmlGridLicensableFeatures_v2_t(_PrintableStructure): ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v2_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), ] + class c_nvmlGridLicensableFeature_t(_PrintableStructure): _fields_ = [ ('featureCode', _nvmlGridLicenseFeatureCode_t), @@ -1751,6 +1830,7 @@ class c_nvmlGridLicensableFeature_t(_PrintableStructure): ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), ] + class c_nvmlGridLicensableFeatures_t(_PrintableStructure): _fields_ = [ ('isGridLicenseSupported', c_int), @@ -1758,257 +1838,239 @@ class c_nvmlGridLicensableFeatures_t(_PrintableStructure): ('gridLicensableFeatures', c_nvmlGridLicensableFeature_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), ] + class c_nvmlMarginTemperature_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), ('marginTemperature', c_int), ] + nvmlMarginTemperature_v1 = 0x1000008 + ## Event structures class struct_c_nvmlEventSet_t(Structure): - pass # opaque handle + pass # opaque handle + + c_nvmlEventSet_t = POINTER(struct_c_nvmlEventSet_t) -nvmlEventTypeSingleBitEccError = 0x0000000000000001 -nvmlEventTypeDoubleBitEccError = 0x0000000000000002 -nvmlEventTypePState = 0x0000000000000004 -nvmlEventTypeXidCriticalError = 0x0000000000000008 -nvmlEventTypeClock = 0x0000000000000010 -nvmlEventTypePowerSourceChange = 0x0000000000000080 -nvmlEventMigConfigChange = 0x0000000000000100 +nvmlEventTypeSingleBitEccError = 0x0000000000000001 +nvmlEventTypeDoubleBitEccError = 0x0000000000000002 +nvmlEventTypePState = 0x0000000000000004 +nvmlEventTypeXidCriticalError = 0x0000000000000008 +nvmlEventTypeClock = 0x0000000000000010 +nvmlEventTypePowerSourceChange = 0x0000000000000080 +nvmlEventMigConfigChange = 0x0000000000000100 nvmlEventTypeSingleBitEccErrorStorm = 0x0000000000000200 -nvmlEventTypeDramRetirementEvent = 0x0000000000000400 -nvmlEventTypeDramRetirementFailure = 0x0000000000000800 -nvmlEventTypeNonFatalPoisonError = 0x0000000000001000 -nvmlEventTypeFatalPoisonError = 0x0000000000002000 -nvmlEventTypeGpuUnavailableError = 0x0000000000004000 -nvmlEventTypeGpuRecoveryAction = 0x0000000000008000 -nvmlEventTypeNone = 0x0000000000000000 -nvmlEventTypeAll = ( - nvmlEventTypeNone - | nvmlEventTypeSingleBitEccError - | nvmlEventTypeDoubleBitEccError - | nvmlEventTypePState - | nvmlEventTypeClock - | nvmlEventTypePowerSourceChange - | nvmlEventTypeXidCriticalError - | nvmlEventMigConfigChange - | nvmlEventTypeSingleBitEccErrorStorm - | nvmlEventTypeDramRetirementEvent - | nvmlEventTypeDramRetirementFailure - | nvmlEventTypeNonFatalPoisonError - | nvmlEventTypeFatalPoisonError - | nvmlEventTypeGpuUnavailableError - | nvmlEventTypeGpuRecoveryAction - ) +nvmlEventTypeDramRetirementEvent = 0x0000000000000400 +nvmlEventTypeDramRetirementFailure = 0x0000000000000800 +nvmlEventTypeNonFatalPoisonError = 0x0000000000001000 +nvmlEventTypeFatalPoisonError = 0x0000000000002000 +nvmlEventTypeGpuUnavailableError = 0x0000000000004000 +nvmlEventTypeGpuRecoveryAction = 0x0000000000008000 +nvmlEventTypeNone = 0x0000000000000000 +nvmlEventTypeAll = (nvmlEventTypeNone + | nvmlEventTypeSingleBitEccError + | nvmlEventTypeDoubleBitEccError + | nvmlEventTypePState + | nvmlEventTypeClock + | nvmlEventTypePowerSourceChange + | nvmlEventTypeXidCriticalError + | nvmlEventMigConfigChange + | nvmlEventTypeSingleBitEccErrorStorm + | nvmlEventTypeDramRetirementEvent + | nvmlEventTypeDramRetirementFailure + | nvmlEventTypeNonFatalPoisonError + | nvmlEventTypeFatalPoisonError + | nvmlEventTypeGpuUnavailableError + | nvmlEventTypeGpuRecoveryAction) ## Clock Event Reasons defines -nvmlClocksEventReasonGpuIdle = 0x0000000000000001 +nvmlClocksEventReasonGpuIdle = 0x0000000000000001 nvmlClocksEventReasonApplicationsClocksSetting = 0x0000000000000002 -nvmlClocksEventReasonUserDefinedClocks = nvmlClocksEventReasonApplicationsClocksSetting # deprecated, use nvmlClocksEventReasonApplicationsClocksSetting -nvmlClocksEventReasonSwPowerCap = 0x0000000000000004 -nvmlClocksEventReasonHwSlowdown = 0x0000000000000008 -nvmlClocksEventReasonSyncBoost = 0x0000000000000010 -nvmlClocksEventReasonSwThermalSlowdown = 0x0000000000000020 -nvmlClocksEventReasonHwThermalSlowdown = 0x0000000000000040 +nvmlClocksEventReasonUserDefinedClocks = nvmlClocksEventReasonApplicationsClocksSetting # deprecated, use nvmlClocksEventReasonApplicationsClocksSetting +nvmlClocksEventReasonSwPowerCap = 0x0000000000000004 +nvmlClocksEventReasonHwSlowdown = 0x0000000000000008 +nvmlClocksEventReasonSyncBoost = 0x0000000000000010 +nvmlClocksEventReasonSwThermalSlowdown = 0x0000000000000020 +nvmlClocksEventReasonHwThermalSlowdown = 0x0000000000000040 nvmlClocksEventReasonHwPowerBrakeSlowdown = 0x0000000000000080 -nvmlClocksEventReasonDisplayClockSetting = 0x0000000000000100 -nvmlClocksEventReasonNone = 0x0000000000000000 -nvmlClocksEventReasonAll = ( - nvmlClocksEventReasonNone | - nvmlClocksEventReasonGpuIdle | - nvmlClocksEventReasonApplicationsClocksSetting | - nvmlClocksEventReasonSwPowerCap | - nvmlClocksEventReasonHwSlowdown | - nvmlClocksEventReasonSyncBoost | - nvmlClocksEventReasonSwThermalSlowdown | - nvmlClocksEventReasonHwThermalSlowdown | - nvmlClocksEventReasonHwPowerBrakeSlowdown | - nvmlClocksEventReasonDisplayClockSetting - ) +nvmlClocksEventReasonDisplayClockSetting = 0x0000000000000100 +nvmlClocksEventReasonNone = 0x0000000000000000 +nvmlClocksEventReasonAll = (nvmlClocksEventReasonNone | nvmlClocksEventReasonGpuIdle + | nvmlClocksEventReasonApplicationsClocksSetting | nvmlClocksEventReasonSwPowerCap + | nvmlClocksEventReasonHwSlowdown | nvmlClocksEventReasonSyncBoost + | nvmlClocksEventReasonSwThermalSlowdown | nvmlClocksEventReasonHwThermalSlowdown + | nvmlClocksEventReasonHwPowerBrakeSlowdown | nvmlClocksEventReasonDisplayClockSetting) ## Following have been deprecated -nvmlClocksThrottleReasonGpuIdle = 0x0000000000000001 +nvmlClocksThrottleReasonGpuIdle = 0x0000000000000001 nvmlClocksThrottleReasonApplicationsClocksSetting = 0x0000000000000002 -nvmlClocksThrottleReasonUserDefinedClocks = nvmlClocksThrottleReasonApplicationsClocksSetting # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting -nvmlClocksThrottleReasonSwPowerCap = 0x0000000000000004 -nvmlClocksThrottleReasonHwSlowdown = 0x0000000000000008 -nvmlClocksThrottleReasonSyncBoost = 0x0000000000000010 -nvmlClocksThrottleReasonSwThermalSlowdown = 0x0000000000000020 -nvmlClocksThrottleReasonHwThermalSlowdown = 0x0000000000000040 +nvmlClocksThrottleReasonUserDefinedClocks = nvmlClocksThrottleReasonApplicationsClocksSetting # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting +nvmlClocksThrottleReasonSwPowerCap = 0x0000000000000004 +nvmlClocksThrottleReasonHwSlowdown = 0x0000000000000008 +nvmlClocksThrottleReasonSyncBoost = 0x0000000000000010 +nvmlClocksThrottleReasonSwThermalSlowdown = 0x0000000000000020 +nvmlClocksThrottleReasonHwThermalSlowdown = 0x0000000000000040 nvmlClocksThrottleReasonHwPowerBrakeSlowdown = 0x0000000000000080 -nvmlClocksThrottleReasonDisplayClockSetting = 0x0000000000000100 -nvmlClocksThrottleReasonNone = 0x0000000000000000 -nvmlClocksThrottleReasonAll = ( - nvmlClocksThrottleReasonNone | - nvmlClocksThrottleReasonGpuIdle | - nvmlClocksThrottleReasonApplicationsClocksSetting | - nvmlClocksThrottleReasonSwPowerCap | - nvmlClocksThrottleReasonHwSlowdown | - nvmlClocksThrottleReasonSyncBoost | - nvmlClocksThrottleReasonSwThermalSlowdown | - nvmlClocksThrottleReasonHwThermalSlowdown | - nvmlClocksThrottleReasonHwPowerBrakeSlowdown | - nvmlClocksThrottleReasonDisplayClockSetting - ) +nvmlClocksThrottleReasonDisplayClockSetting = 0x0000000000000100 +nvmlClocksThrottleReasonNone = 0x0000000000000000 +nvmlClocksThrottleReasonAll = (nvmlClocksThrottleReasonNone | nvmlClocksThrottleReasonGpuIdle + | nvmlClocksThrottleReasonApplicationsClocksSetting | nvmlClocksThrottleReasonSwPowerCap + | nvmlClocksThrottleReasonHwSlowdown | nvmlClocksThrottleReasonSyncBoost + | nvmlClocksThrottleReasonSwThermalSlowdown | nvmlClocksThrottleReasonHwThermalSlowdown + | nvmlClocksThrottleReasonHwPowerBrakeSlowdown + | nvmlClocksThrottleReasonDisplayClockSetting) + class c_nvmlEventData_t(_PrintableStructure): - _fields_ = [ - ('device', c_nvmlDevice_t), - ('eventType', c_ulonglong), - ('eventData', c_ulonglong), - ('gpuInstanceId', c_uint), - ('computeInstanceId', c_uint) - ] + _fields_ = [('device', c_nvmlDevice_t), ('eventType', c_ulonglong), ('eventData', c_ulonglong), + ('gpuInstanceId', c_uint), ('computeInstanceId', c_uint)] _fmt_ = {'eventType': "0x%08X"} + class c_nvmlAccountingStats_t(_PrintableStructure): - _fields_ = [ - ('gpuUtilization', c_uint), - ('memoryUtilization', c_uint), - ('maxMemoryUsage', c_ulonglong), - ('time', c_ulonglong), - ('startTime', c_ulonglong), - ('isRunning', c_uint), - ('reserved', c_uint * 5) - ] + _fields_ = [('gpuUtilization', c_uint), ('memoryUtilization', c_uint), ('maxMemoryUsage', c_ulonglong), + ('time', c_ulonglong), ('startTime', c_ulonglong), ('isRunning', c_uint), ('reserved', c_uint * 5)] + class c_nvmlVgpuVersion_t(Structure): - _fields_ = [("minVersion", c_uint), - ("maxVersion", c_uint) - ] + _fields_ = [("minVersion", c_uint), ("maxVersion", c_uint)] + class c_nvmlVgpuMetadata_t(_PrintableStructure): - _fields_ = [("version", c_uint), - ("revision", c_uint), - ("guestInfoState", _nvmlVgpuGuestInfoState_t), + _fields_ = [("version", c_uint), ("revision", c_uint), ("guestInfoState", _nvmlVgpuGuestInfoState_t), ("guestDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), - ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), - ("reserved", c_uint * 6), - ("vgpuVirtualizationCaps", c_uint), - ("guestVgpuVersion", c_uint), - ("opaqueDataSize", c_uint), - ("opaqueData", c_char * NVML_VGPU_METADATA_OPAQUE_DATA_SIZE) - ] + ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), ("reserved", c_uint * 6), + ("vgpuVirtualizationCaps", c_uint), ("guestVgpuVersion", c_uint), ("opaqueDataSize", c_uint), + ("opaqueData", c_char * NVML_VGPU_METADATA_OPAQUE_DATA_SIZE)] + class c_nvmlVgpuPgpuMetadata_t(_PrintableStructure): - _fields_ = [("version", c_uint), - ("revision", c_uint), + _fields_ = [("version", c_uint), ("revision", c_uint), ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), - ("pgpuVirtualizationCaps", c_uint), - ("reserved", c_uint * 5), - ("hostSupportedVgpuRange", c_nvmlVgpuVersion_t), - ("opaqueDataSize", c_uint), - ("opaqueData", c_char * NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE) - ] + ("pgpuVirtualizationCaps", c_uint), ("reserved", c_uint * 5), + ("hostSupportedVgpuRange", c_nvmlVgpuVersion_t), ("opaqueDataSize", c_uint), + ("opaqueData", c_char * NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE)] + class c_nvmlVgpuPgpuCompatibility_t(Structure): _fields_ = [("vgpuVmCompatibility", _nvmlVgpuVmCompatibility_t), - ("compatibilityLimitCode", _nvmlVgpuPgpuCompatibilityLimitCode_t) - ] + ("compatibilityLimitCode", _nvmlVgpuPgpuCompatibilityLimitCode_t)] + ## vGPU scheduler policy defines -NVML_VGPU_SCHEDULER_POLICY_UNKNOWN = 0 -NVML_VGPU_SCHEDULER_POLICY_BEST_EFFORT = 1 -NVML_VGPU_SCHEDULER_POLICY_EQUAL_SHARE = 2 -NVML_VGPU_SCHEDULER_POLICY_FIXED_SHARE = 3 +NVML_VGPU_SCHEDULER_POLICY_UNKNOWN = 0 +NVML_VGPU_SCHEDULER_POLICY_BEST_EFFORT = 1 +NVML_VGPU_SCHEDULER_POLICY_EQUAL_SHARE = 2 +NVML_VGPU_SCHEDULER_POLICY_FIXED_SHARE = 3 ## Supported vGPU scheduler policy count -NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT = 3 +NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT = 3 -NVML_SCHEDULER_SW_MAX_LOG_ENTRIES = 200 +NVML_SCHEDULER_SW_MAX_LOG_ENTRIES = 200 + +NVML_VGPU_SCHEDULER_ARR_DEFAULT = 0 +NVML_VGPU_SCHEDULER_ARR_DISABLE = 1 +NVML_VGPU_SCHEDULER_ARR_ENABLE = 2 -NVML_VGPU_SCHEDULER_ARR_DEFAULT = 0 -NVML_VGPU_SCHEDULER_ARR_DISABLE = 1 -NVML_VGPU_SCHEDULER_ARR_ENABLE = 2 class c_nvmlVgpuSchedDataWithARR_t(_PrintableStructure): _fields_ = [ - ('avgFactor', c_uint), - ('timeslice', c_uint), + ('avgFactor', c_uint), + ('timeslice', c_uint), ] + class c_nvmlVgpuSchedData_t(_PrintableStructure): _fields_ = [ - ('timeslice', c_uint), + ('timeslice', c_uint), ] + class c_nvmlVgpuSchedulerParams_t(Union): _fields_ = [ ('vgpuSchedDataWithARR', c_nvmlVgpuSchedDataWithARR_t), - ('vgpuSchedData', c_nvmlVgpuSchedData_t), + ('vgpuSchedData', c_nvmlVgpuSchedData_t), ] + class c_nvmlVgpuSchedulerLogEntry_t(_PrintableStructure): _fields_ = [ - ('timestamp', c_ulonglong), - ('timeRunTotal', c_ulonglong), - ('timeRun', c_ulonglong), - ('swRunlistId', c_uint), - ('targetTimeSlice', c_ulonglong), - ('cumulativePreemptionTime', c_ulonglong), + ('timestamp', c_ulonglong), + ('timeRunTotal', c_ulonglong), + ('timeRun', c_ulonglong), + ('swRunlistId', c_uint), + ('targetTimeSlice', c_ulonglong), + ('cumulativePreemptionTime', c_ulonglong), ] + class c_nvmlVgpuSchedulerLog_t(_PrintableStructure): _fields_ = [ - ('engineId', c_uint), + ('engineId', c_uint), ('schedulerPolicy', c_uint), - ('arrMode', c_uint), + ('arrMode', c_uint), ('schedulerParams', c_nvmlVgpuSchedulerParams_t), - ('entriesCount', c_uint), - ('logEntries', c_nvmlVgpuSchedulerLogEntry_t * NVML_SCHEDULER_SW_MAX_LOG_ENTRIES), + ('entriesCount', c_uint), + ('logEntries', c_nvmlVgpuSchedulerLogEntry_t * NVML_SCHEDULER_SW_MAX_LOG_ENTRIES), ] + class c_nvmlVgpuSchedulerGetState_t(_PrintableStructure): _fields_ = [ ('schedulerPolicy', c_uint), - ('arrMode', c_uint), + ('arrMode', c_uint), ('schedulerParams', c_nvmlVgpuSchedulerParams_t), ] + class c_nvmlVgpuSchedSetDataWithARR_t(_PrintableStructure): _fields_ = [ - ('avgFactor', c_uint), - ('frequency', c_uint), + ('avgFactor', c_uint), + ('frequency', c_uint), ] + class c_nvmlVgpuSchedSetData_t(_PrintableStructure): _fields_ = [ - ('timeslice', c_uint), + ('timeslice', c_uint), ] + class c_nvmlVgpuSchedulerSetParams_t(Union): _fields_ = [ ('vgpuSchedDataWithARR', c_nvmlVgpuSchedSetDataWithARR_t), - ('vgpuSchedData', c_nvmlVgpuSchedSetData_t), + ('vgpuSchedData', c_nvmlVgpuSchedSetData_t), ] + class c_nvmlVgpuSchedulerSetState_t(_PrintableStructure): _fields_ = [ ('schedulerPolicy', c_uint), - ('enableARRMode', c_uint), + ('enableARRMode', c_uint), ('schedulerParams', c_nvmlVgpuSchedulerSetParams_t), ] + class c_nvmlVgpuSchedulerCapabilities_t(_PrintableStructure): _fields_ = [ ('supportedSchedulers', c_uint * NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT), - ('maxTimeslice', c_uint), - ('minTimeslice', c_uint), - ('isArrModeSupported', c_uint), - ('maxFrequencyForARR', c_uint), - ('minFrequencyForARR', c_uint), - ('maxAvgFactorForARR', c_uint), - ('minAvgFactorForARR', c_uint), + ('maxTimeslice', c_uint), + ('minTimeslice', c_uint), + ('isArrModeSupported', c_uint), + ('maxFrequencyForARR', c_uint), + ('minFrequencyForARR', c_uint), + ('maxAvgFactorForARR', c_uint), + ('minAvgFactorForARR', c_uint), ] + class c_nvmlFBCStats_t(Structure): - _fields_ = [("sessionsCount", c_uint), - ("averageFPS", c_uint), - ("averageLatency", c_uint) - ] + _fields_ = [("sessionsCount", c_uint), ("averageFPS", c_uint), ("averageLatency", c_uint)] + class c_nvmlFBCSession_t(_PrintableStructure): _fields_ = [ @@ -2026,211 +2088,193 @@ class c_nvmlFBCSession_t(_PrintableStructure): ('averageLatency', c_uint), ] -NVML_DEVICE_MIG_DISABLE = 0x0 -NVML_DEVICE_MIG_ENABLE = 0x1 -NVML_GPU_INSTANCE_PROFILE_1_SLICE = 0x0 -NVML_GPU_INSTANCE_PROFILE_2_SLICE = 0x1 -NVML_GPU_INSTANCE_PROFILE_3_SLICE = 0x2 -NVML_GPU_INSTANCE_PROFILE_4_SLICE = 0x3 -NVML_GPU_INSTANCE_PROFILE_7_SLICE = 0x4 -NVML_GPU_INSTANCE_PROFILE_8_SLICE = 0x5 -NVML_GPU_INSTANCE_PROFILE_6_SLICE = 0x6 +NVML_DEVICE_MIG_DISABLE = 0x0 +NVML_DEVICE_MIG_ENABLE = 0x1 + +NVML_GPU_INSTANCE_PROFILE_1_SLICE = 0x0 +NVML_GPU_INSTANCE_PROFILE_2_SLICE = 0x1 +NVML_GPU_INSTANCE_PROFILE_3_SLICE = 0x2 +NVML_GPU_INSTANCE_PROFILE_4_SLICE = 0x3 +NVML_GPU_INSTANCE_PROFILE_7_SLICE = 0x4 +NVML_GPU_INSTANCE_PROFILE_8_SLICE = 0x5 +NVML_GPU_INSTANCE_PROFILE_6_SLICE = 0x6 NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 NVML_GPU_INSTANCE_PROFILE_2_SLICE_REV1 = 0x8 NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV2 = 0x9 -NVML_GPU_INSTANCE_PROFILE_1_SLICE_GFX = 0xA -NVML_GPU_INSTANCE_PROFILE_2_SLICE_GFX = 0xB -NVML_GPU_INSTANCE_PROFILE_4_SLICE_GFX = 0xC -NVML_GPU_INSTANCE_PROFILE_COUNT = 0xD +NVML_GPU_INSTANCE_PROFILE_1_SLICE_GFX = 0xA +NVML_GPU_INSTANCE_PROFILE_2_SLICE_GFX = 0xB +NVML_GPU_INSTANCE_PROFILE_4_SLICE_GFX = 0xC +NVML_GPU_INSTANCE_PROFILE_COUNT = 0xD + class c_nvmlGpuInstancePlacement_t(Structure): - _fields_ = [("start", c_uint), - ("size", c_uint) - ] + _fields_ = [("start", c_uint), ("size", c_uint)] + class c_nvmlGpuInstanceProfileInfo_t(Structure): - _fields_ = [("id", c_uint), - ("isP2pSupported", c_uint), - ("sliceCount", c_uint), - ("instanceCount", c_uint), - ("multiprocessorCount", c_uint), - ("copyEngineCount", c_uint), - ("decoderCount", c_uint), - ("encoderCount", c_uint), - ("jpegCount", c_uint), - ("ofaCount", c_uint), - ("memorySizeMB", c_ulonglong), - ] + _fields_ = [ + ("id", c_uint), + ("isP2pSupported", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("copyEngineCount", c_uint), + ("decoderCount", c_uint), + ("encoderCount", c_uint), + ("jpegCount", c_uint), + ("ofaCount", c_uint), + ("memorySizeMB", c_ulonglong), + ] + nvmlGpuInstanceProfileInfo_v2 = 0x02000098 + class c_nvmlGpuInstanceProfileInfo_v2_t(_PrintableStructure): - _fields_ = [("version", c_uint), - ("id", c_uint), - ("isP2pSupported", c_uint), - ("sliceCount", c_uint), - ("instanceCount", c_uint), - ("multiprocessorCount", c_uint), - ("copyEngineCount", c_uint), - ("decoderCount", c_uint), - ("encoderCount", c_uint), - ("jpegCount", c_uint), - ("ofaCount", c_uint), - ("memorySizeMB", c_ulonglong), - ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE) - ] - + _fields_ = [("version", c_uint), ("id", c_uint), ("isP2pSupported", c_uint), ("sliceCount", c_uint), + ("instanceCount", c_uint), ("multiprocessorCount", c_uint), ("copyEngineCount", c_uint), + ("decoderCount", c_uint), ("encoderCount", c_uint), ("jpegCount", c_uint), ("ofaCount", c_uint), + ("memorySizeMB", c_ulonglong), ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE)] + def __init__(self): super(c_nvmlGpuInstanceProfileInfo_v2_t, self).__init__(version=nvmlGpuInstanceProfileInfo_v2) + class c_nvmlGpuInstanceInfo_t(Structure): - _fields_ = [("device", c_nvmlDevice_t), - ("id", c_uint), - ("profileId", c_uint), - ("placement", c_nvmlGpuInstancePlacement_t) - ] + _fields_ = [("device", c_nvmlDevice_t), ("id", c_uint), ("profileId", c_uint), + ("placement", c_nvmlGpuInstancePlacement_t)] + class struct_c_nvmlGpuInstance_t(Structure): - pass # opaque handle + pass # opaque handle + + c_nvmlGpuInstance_t = POINTER(struct_c_nvmlGpuInstance_t) -NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE = 0x0 -NVML_COMPUTE_INSTANCE_PROFILE_2_SLICE = 0x1 -NVML_COMPUTE_INSTANCE_PROFILE_3_SLICE = 0x2 -NVML_COMPUTE_INSTANCE_PROFILE_4_SLICE = 0x3 -NVML_COMPUTE_INSTANCE_PROFILE_7_SLICE = 0x4 -NVML_COMPUTE_INSTANCE_PROFILE_8_SLICE = 0x5 -NVML_COMPUTE_INSTANCE_PROFILE_6_SLICE = 0x6 +NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE = 0x0 +NVML_COMPUTE_INSTANCE_PROFILE_2_SLICE = 0x1 +NVML_COMPUTE_INSTANCE_PROFILE_3_SLICE = 0x2 +NVML_COMPUTE_INSTANCE_PROFILE_4_SLICE = 0x3 +NVML_COMPUTE_INSTANCE_PROFILE_7_SLICE = 0x4 +NVML_COMPUTE_INSTANCE_PROFILE_8_SLICE = 0x5 +NVML_COMPUTE_INSTANCE_PROFILE_6_SLICE = 0x6 NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 -NVML_COMPUTE_INSTANCE_PROFILE_COUNT = 0x8 +NVML_COMPUTE_INSTANCE_PROFILE_COUNT = 0x8 NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED = 0x0 NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT = 0x1 + class c_nvmlComputeInstancePlacement_t(Structure): - _fields_ = [("start", c_uint), - ("size", c_uint) - ] + _fields_ = [("start", c_uint), ("size", c_uint)] + class c_nvmlComputeInstanceProfileInfo_t(Structure): - _fields_ = [("id", c_uint), - ("sliceCount", c_uint), - ("instanceCount", c_uint), - ("multiprocessorCount", c_uint), - ("sharedCopyEngineCount", c_uint), - ("sharedDecoderCount", c_uint), - ("sharedEncoderCount", c_uint), - ("sharedJpegCount", c_uint), - ("sharedOfaCount", c_uint) - ] + _fields_ = [("id", c_uint), ("sliceCount", c_uint), ("instanceCount", c_uint), ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), ("sharedDecoderCount", c_uint), ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), ("sharedOfaCount", c_uint)] + nvmlComputeInstanceProfileInfo_v2 = 0x02000088 + class c_nvmlComputeInstanceProfileInfo_v2_t(_PrintableStructure): - _fields_ = [("version", c_uint), - ("id", c_uint), - ("sliceCount", c_uint), - ("instanceCount", c_uint), - ("multiprocessorCount", c_uint), - ("sharedCopyEngineCount", c_uint), - ("sharedDecoderCount", c_uint), - ("sharedEncoderCount", c_uint), - ("sharedJpegCount", c_uint), - ("sharedOfaCount", c_uint), - ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE) - ] + _fields_ = [("version", c_uint), ("id", c_uint), ("sliceCount", c_uint), ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), ("sharedCopyEngineCount", c_uint), ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), ("sharedJpegCount", c_uint), ("sharedOfaCount", c_uint), + ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE)] def __init__(self): super(c_nvmlComputeInstanceProfileInfo_v2_t, self).__init__(version=nvmlComputeInstanceProfileInfo_v2) + class c_nvmlComputeInstanceInfo_t(Structure): - _fields_ = [("device", c_nvmlDevice_t), - ("gpuInstance", c_nvmlGpuInstance_t), - ("id", c_uint), - ("profileId", c_uint), - ("placement", c_nvmlComputeInstancePlacement_t) - ] + _fields_ = [("device", c_nvmlDevice_t), ("gpuInstance", c_nvmlGpuInstance_t), ("id", c_uint), ("profileId", c_uint), + ("placement", c_nvmlComputeInstancePlacement_t)] + NVML_MAX_GPU_UTILIZATIONS = 8 -NVML_GPU_UTILIZATION_DOMAIN_GPU = 0 -NVML_GPU_UTILIZATION_DOMAIN_FB = 1 -NVML_GPU_UTILIZATION_DOMAIN_VID = 2 -NVML_GPU_UTILIZATION_DOMAIN_BUS = 3 +NVML_GPU_UTILIZATION_DOMAIN_GPU = 0 +NVML_GPU_UTILIZATION_DOMAIN_FB = 1 +NVML_GPU_UTILIZATION_DOMAIN_VID = 2 +NVML_GPU_UTILIZATION_DOMAIN_BUS = 3 + + class c_nvmlGpuDynamicPstatesUtilization_t(Structure): - _fields_ = [("bIsPresent", c_uint, 1), - ("percentage", c_uint), - ("incThreshold", c_uint), - ("decThreshold", c_uint)] + _fields_ = [("bIsPresent", c_uint, 1), ("percentage", c_uint), ("incThreshold", c_uint), ("decThreshold", c_uint)] + + class c_nvmlGpuDynamicPstatesInfo_t(Structure): - _fields_ = [("flags", c_uint), - ("utilization", c_nvmlGpuDynamicPstatesUtilization_t * NVML_MAX_GPU_UTILIZATIONS)] + _fields_ = [("flags", c_uint), ("utilization", c_nvmlGpuDynamicPstatesUtilization_t * NVML_MAX_GPU_UTILIZATIONS)] + NVML_MAX_THERMAL_SENSORS_PER_GPU = 3 -NVML_THERMAL_TARGET_NONE = 0 -NVML_THERMAL_TARGET_GPU = 1 -NVML_THERMAL_TARGET_MEMORY = 2 -NVML_THERMAL_TARGET_POWER_SUPPLY = 4 -NVML_THERMAL_TARGET_BOARD = 8 -NVML_THERMAL_TARGET_VCD_BOARD = 9 -NVML_THERMAL_TARGET_VCD_INLET = 10 -NVML_THERMAL_TARGET_VCD_OUTLET = 11 -NVML_THERMAL_TARGET_ALL = 15 -NVML_THERMAL_TARGET_UNKNOWN = -1 +NVML_THERMAL_TARGET_NONE = 0 +NVML_THERMAL_TARGET_GPU = 1 +NVML_THERMAL_TARGET_MEMORY = 2 +NVML_THERMAL_TARGET_POWER_SUPPLY = 4 +NVML_THERMAL_TARGET_BOARD = 8 +NVML_THERMAL_TARGET_VCD_BOARD = 9 +NVML_THERMAL_TARGET_VCD_INLET = 10 +NVML_THERMAL_TARGET_VCD_OUTLET = 11 +NVML_THERMAL_TARGET_ALL = 15 +NVML_THERMAL_TARGET_UNKNOWN = -1 -NVML_THERMAL_CONTROLLER_NONE = 0 -NVML_THERMAL_CONTROLLER_GPU_INTERNAL = 1 -NVML_THERMAL_CONTROLLER_ADM1032 = 2 -NVML_THERMAL_CONTROLLER_ADT7461 = 3 -NVML_THERMAL_CONTROLLER_MAX6649 = 4 -NVML_THERMAL_CONTROLLER_MAX1617 = 5 -NVML_THERMAL_CONTROLLER_LM99 = 6 -NVML_THERMAL_CONTROLLER_LM89 = 7 -NVML_THERMAL_CONTROLLER_LM64 = 8 -NVML_THERMAL_CONTROLLER_G781 = 9 -NVML_THERMAL_CONTROLLER_ADT7473 = 10 -NVML_THERMAL_CONTROLLER_SBMAX6649 = 11 -NVML_THERMAL_CONTROLLER_VBIOSEVT = 12 -NVML_THERMAL_CONTROLLER_OS = 13 +NVML_THERMAL_CONTROLLER_NONE = 0 +NVML_THERMAL_CONTROLLER_GPU_INTERNAL = 1 +NVML_THERMAL_CONTROLLER_ADM1032 = 2 +NVML_THERMAL_CONTROLLER_ADT7461 = 3 +NVML_THERMAL_CONTROLLER_MAX6649 = 4 +NVML_THERMAL_CONTROLLER_MAX1617 = 5 +NVML_THERMAL_CONTROLLER_LM99 = 6 +NVML_THERMAL_CONTROLLER_LM89 = 7 +NVML_THERMAL_CONTROLLER_LM64 = 8 +NVML_THERMAL_CONTROLLER_G781 = 9 +NVML_THERMAL_CONTROLLER_ADT7473 = 10 +NVML_THERMAL_CONTROLLER_SBMAX6649 = 11 +NVML_THERMAL_CONTROLLER_VBIOSEVT = 12 +NVML_THERMAL_CONTROLLER_OS = 13 NVML_THERMAL_CONTROLLER_NVSYSCON_CANOAS = 14 -NVML_THERMAL_CONTROLLER_NVSYSCON_E551 = 15 -NVML_THERMAL_CONTROLLER_MAX6649R = 16 -NVML_THERMAL_CONTROLLER_ADT7473S = 17 -NVML_THERMAL_CONTROLLER_UNKNOWN = -1 +NVML_THERMAL_CONTROLLER_NVSYSCON_E551 = 15 +NVML_THERMAL_CONTROLLER_MAX6649R = 16 +NVML_THERMAL_CONTROLLER_ADT7473S = 17 +NVML_THERMAL_CONTROLLER_UNKNOWN = -1 + class c_nvmlGpuThermalSensor_t(Structure): - _fields_ = [("controller", c_int), - ("defaultMinTemp", c_int), - ("defaultMaxTemp", c_int), - ("currentTemp", c_int), + _fields_ = [("controller", c_int), ("defaultMinTemp", c_int), ("defaultMaxTemp", c_int), ("currentTemp", c_int), ("target", c_int)] + + class c_nvmlGpuThermalSettings_t(Structure): - _fields_ = [("count", c_uint), - ("sensor", c_nvmlGpuThermalSensor_t * NVML_MAX_THERMAL_SENSORS_PER_GPU)] + _fields_ = [("count", c_uint), ("sensor", c_nvmlGpuThermalSensor_t * NVML_MAX_THERMAL_SENSORS_PER_GPU)] + _nvmlCoolerControl_t = c_uint -NVML_THERMAL_COOLER_SIGNAL_NONE = 0 -NVML_THERMAL_COOLER_SIGNAL_TOGGLE = 1 -NVML_THERMAL_COOLER_SIGNAL_VARIABLE = 2 -NVML_THERMAL_COOLER_SIGNAL_COUNT = 3 +NVML_THERMAL_COOLER_SIGNAL_NONE = 0 +NVML_THERMAL_COOLER_SIGNAL_TOGGLE = 1 +NVML_THERMAL_COOLER_SIGNAL_VARIABLE = 2 +NVML_THERMAL_COOLER_SIGNAL_COUNT = 3 _nvmlCoolerTarget_t = c_uint -NVML_THERMAL_COOLER_TARGET_NONE = (1 << 0) -NVML_THERMAL_COOLER_TARGET_GPU = (1 << 1) -NVML_THERMAL_COOLER_TARGET_MEMORY = (1 << 2) -NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY = (1 << 3) -NVML_THERMAL_COOLER_TARGET_GPU_RELATED = (NVML_THERMAL_COOLER_TARGET_GPU | NVML_THERMAL_COOLER_TARGET_MEMORY | NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY) +NVML_THERMAL_COOLER_TARGET_NONE = (1 << 0) +NVML_THERMAL_COOLER_TARGET_GPU = (1 << 1) +NVML_THERMAL_COOLER_TARGET_MEMORY = (1 << 2) +NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY = (1 << 3) +NVML_THERMAL_COOLER_TARGET_GPU_RELATED = (NVML_THERMAL_COOLER_TARGET_GPU | NVML_THERMAL_COOLER_TARGET_MEMORY + | NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY) + class c_nvmlCoolerInfo_t(_PrintableStructure): - _fields_ = [("version", c_uint), - ("index", c_uint), - ("coolerControlType", _nvmlCoolerControl_t), - ("coolerTarget", _nvmlCoolerTarget_t) - ] + _fields_ = [("version", c_uint), ("index", c_uint), ("coolerControlType", _nvmlCoolerControl_t), + ("coolerTarget", _nvmlCoolerTarget_t)] + nvmlCoolerInfo_v1 = 0x1000010 + def nvmlDeviceGetCoolerInfo(handle): c_coolerInfo = c_nvmlCoolerInfo_t() c_coolerInfo.version = nvmlCoolerInfo_v1 @@ -2240,92 +2284,116 @@ def nvmlDeviceGetCoolerInfo(handle): _nvmlCheckReturn(ret) return [c_coolerInfo.coolerControlType, c_coolerInfo.coolerTarget] + class struct_c_nvmlComputeInstance_t(Structure): - pass # opaque handle + pass # opaque handle + + c_nvmlComputeInstance_t = POINTER(struct_c_nvmlComputeInstance_t) + class c_nvmlDeviceAttributes(Structure): - _fields_ = [("multiprocessorCount", c_uint), - ("sharedCopyEngineCount", c_uint), - ("sharedDecoderCount", c_uint), - ("sharedEncoderCount", c_uint), - ("sharedJpegCount", c_uint), - ("sharedOfaCount", c_uint), - ("gpuInstanceSliceCount", c_uint), - ("computeInstanceSliceCount", c_uint), - ("memorySizeMB", c_ulonglong), - ] + _fields_ = [ + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ("gpuInstanceSliceCount", c_uint), + ("computeInstanceSliceCount", c_uint), + ("memorySizeMB", c_ulonglong), + ] + class c_nvmlRowRemapperHistogramValues(Structure): - _fields_ = [("max", c_uint), - ("high", c_uint), - ("partial", c_uint), - ("low", c_uint), - ("none", c_uint) - ] + _fields_ = [("max", c_uint), ("high", c_uint), ("partial", c_uint), ("low", c_uint), ("none", c_uint)] -NVML_GPU_CERT_CHAIN_SIZE = 0x1000 -NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE = 0x1400 -NVML_CC_GPU_CEC_NONCE_SIZE = 0x20 -NVML_CC_GPU_ATTESTATION_REPORT_SIZE = 0x2000 + +NVML_GPU_CERT_CHAIN_SIZE = 0x1000 +NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE = 0x1400 +NVML_CC_GPU_CEC_NONCE_SIZE = 0x20 +NVML_CC_GPU_ATTESTATION_REPORT_SIZE = 0x2000 NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE = 0x1000 NVML_CC_CEC_ATTESTATION_REPORT_NOT_PRESENT = 0 -NVML_CC_CEC_ATTESTATION_REPORT_PRESENT = 1 +NVML_CC_CEC_ATTESTATION_REPORT_PRESENT = 1 + class c_nvmlConfComputeSystemState_t(Structure): - _fields_ = [('environment', c_uint), - ('ccFeature', c_uint), - ('devToolsMode', c_uint), - ] + _fields_ = [ + ('environment', c_uint), + ('ccFeature', c_uint), + ('devToolsMode', c_uint), + ] + nvmlSystemConfComputeSettings_v1 = 0x1000014 + class c_nvmlSystemConfComputeSettings_v1_t(Structure): - _fields_ = [('version', c_uint), - ('environment', c_uint), - ('ccFeature', c_uint), - ('devToolsMode', c_uint), - ('multiGpuMode', c_uint), - ] + _fields_ = [ + ('version', c_uint), + ('environment', c_uint), + ('ccFeature', c_uint), + ('devToolsMode', c_uint), + ('multiGpuMode', c_uint), + ] + def __init__(self): super(c_nvmlSystemConfComputeSettings_v1_t, self).__init__(version=nvmlSystemConfComputeSettings_v1) + class c_nvmlConfComputeSystemCaps_t(Structure): - _fields_ = [('cpuCaps', c_uint), - ('gpusCaps', c_uint), - ] + _fields_ = [ + ('cpuCaps', c_uint), + ('gpusCaps', c_uint), + ] + class c_nvmlConfComputeMemSizeInfo_t(Structure): - _fields_ = [('protectedMemSizeKib', c_ulonglong), - ('unprotectedMemSizeKib', c_ulonglong), - ] + _fields_ = [ + ('protectedMemSizeKib', c_ulonglong), + ('unprotectedMemSizeKib', c_ulonglong), + ] + class c_nvmlConfComputeGpuCertificate_t(Structure): - _fields_ = [('certChainSize', c_uint), - ('attestationCertChainSize', c_uint), - ('certChain', c_uint8 * NVML_GPU_CERT_CHAIN_SIZE), - ('attestationCertChain', c_uint8 * NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE), - ] + _fields_ = [ + ('certChainSize', c_uint), + ('attestationCertChainSize', c_uint), + ('certChain', c_uint8 * NVML_GPU_CERT_CHAIN_SIZE), + ('attestationCertChain', c_uint8 * NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE), + ] + class c_nvmlConfComputeGpuAttestationReport_t(Structure): - _fields_ = [('isCecAttestationReportPresent', c_uint), - ('attestationReportSize', c_uint), - ('cecAttestationReportSize', c_uint), - ('nonce', c_uint8 * NVML_CC_GPU_CEC_NONCE_SIZE), - ('attestationReport', c_uint8 * NVML_CC_GPU_ATTESTATION_REPORT_SIZE), - ('cecAttestationReport', c_uint8 * NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE), - ] + _fields_ = [ + ('isCecAttestationReportPresent', c_uint), + ('attestationReportSize', c_uint), + ('cecAttestationReportSize', c_uint), + ('nonce', c_uint8 * NVML_CC_GPU_CEC_NONCE_SIZE), + ('attestationReport', c_uint8 * NVML_CC_GPU_ATTESTATION_REPORT_SIZE), + ('cecAttestationReport', c_uint8 * NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE), + ] + class c_nvmlConfComputeSetKeyRotationThresholdInfo_t(Structure): - _fields_ = [('version', c_uint), - ('maxAttackerAdvantage', c_ulong), - ] + _fields_ = [ + ('version', c_uint), + ('maxAttackerAdvantage', c_ulong), + ] + + ConfComputeSetKeyRotationThresholdInfo_v1 = 0x1000010 + class c_nvmlConfComputeGetKeyRotationThresholdInfo_t(Structure): - _fields_ = [('version', c_uint), - ('attackerAdvantage', c_ulong), - ] + _fields_ = [ + ('version', c_uint), + ('attackerAdvantage', c_ulong), + ] + + ConfComputeGetKeyRotationThresholdInfo_v1 = 0x1000010 @@ -2339,6 +2407,7 @@ def convertStrBytes(func): Returned from function: b'returned string' Returned to caller: 'returned string' ''' + @wraps(func) def wrapper(*args, **kwargs): # encoding a str returns bytes in python 2 and 3 @@ -2353,11 +2422,13 @@ def convertStrBytes(func): return res.decode() return res - if sys.version_info >= (3,): + if sys.version_info >= (3, ): return wrapper return func + def throwOnVersionMismatch(func): + @wraps(func) def wrapper(*args, **kwargs): try: @@ -2367,8 +2438,10 @@ def throwOnVersionMismatch(func): "pyNVML version does not match the NVML lib version. " "Either use matching pyNVML and NVML lib versions or " "use a versioned function such as " + func.__name__ + "_v2") + return wrapper + ## C function wrappers ## def nvmlInitWithFlags(flags): _LoadNvmlLibrary() @@ -2387,10 +2460,12 @@ def nvmlInitWithFlags(flags): libLoadLock.release() return None + def nvmlInit(): nvmlInitWithFlags(0) return None + def _LoadNvmlLibrary(): ''' Load the library if it isn't loaded already @@ -2413,7 +2488,9 @@ def _LoadNvmlLibrary(): except OSError as ose: # If nvml.dll is not found in System32, it should be in ProgramFiles # load nvml.dll from %ProgramFiles%/NVIDIA Corporation/NVSMI/nvml.dll - nvmlLib = CDLL(os.path.join(os.getenv("ProgramFiles", "C:/Program Files"), "NVIDIA Corporation/NVSMI/nvml.dll")) + nvmlLib = CDLL( + os.path.join(os.getenv("ProgramFiles", "C:/Program Files"), + "NVIDIA Corporation/NVSMI/nvml.dll")) else: # assume linux nvmlLib = CDLL("libnvidia-ml.so.1") @@ -2425,6 +2502,7 @@ def _LoadNvmlLibrary(): # lock is always freed libLoadLock.release() + def nvmlShutdown(): # # Leave the library loaded, but shutdown the interface @@ -2441,14 +2519,16 @@ def nvmlShutdown(): libLoadLock.release() return None + # Added in 2.285 @convertStrBytes def nvmlErrorString(result): fn = _nvmlGetFunctionPointer("nvmlErrorString") - fn.restype = c_char_p # otherwise return is an int + fn.restype = c_char_p # otherwise return is an int ret = fn(result) return ret + # Added in 2.285 @convertStrBytes def nvmlSystemGetNVMLVersion(): @@ -2458,6 +2538,7 @@ def nvmlSystemGetNVMLVersion(): _nvmlCheckReturn(ret) return c_version.value + def nvmlSystemGetCudaDriverVersion(): c_cuda_version = c_int() fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion") @@ -2465,6 +2546,7 @@ def nvmlSystemGetCudaDriverVersion(): _nvmlCheckReturn(ret) return c_cuda_version.value + def nvmlSystemGetCudaDriverVersion_v2(): c_cuda_version = c_int() fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion_v2") @@ -2472,6 +2554,7 @@ def nvmlSystemGetCudaDriverVersion_v2(): _nvmlCheckReturn(ret) return c_cuda_version.value + # Added in 2.285 @convertStrBytes def nvmlSystemGetProcessName(pid): @@ -2481,6 +2564,7 @@ def nvmlSystemGetProcessName(pid): _nvmlCheckReturn(ret) return c_name.value + @convertStrBytes def nvmlSystemGetDriverVersion(): c_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) @@ -2489,6 +2573,7 @@ def nvmlSystemGetDriverVersion(): _nvmlCheckReturn(ret) return c_version.value + # Added in 2.285 def nvmlSystemGetHicVersion(): c_count = c_uint(0) @@ -2499,8 +2584,7 @@ def nvmlSystemGetHicVersion(): ret = fn(byref(c_count), None) # this should only fail with insufficient size - if ((ret != NVML_SUCCESS) and - (ret != NVML_ERROR_INSUFFICIENT_SIZE)): + if ((ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE)): raise NVMLError(ret) # If there are no hics @@ -2513,14 +2597,16 @@ def nvmlSystemGetHicVersion(): _nvmlCheckReturn(ret) return hics + def nvmlSystemGetDriverBranch(): c_branchInfo = c_nvmlSystemDriverBranchInfo_v1_t(0) c_branchInfo.version = SystemDriverBranchInfo_v1 - fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverBranch") + fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverBranch") ret = fn(byref(c_branchInfo), c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_branchInfo + ## Unit get functions def nvmlUnitGetCount(): c_count = c_uint() @@ -2529,6 +2615,7 @@ def nvmlUnitGetCount(): _nvmlCheckReturn(ret) return c_count.value + def nvmlUnitGetHandleByIndex(index): c_index = c_uint(index) unit = c_nvmlUnit_t() @@ -2537,6 +2624,7 @@ def nvmlUnitGetHandleByIndex(index): _nvmlCheckReturn(ret) return unit + def nvmlUnitGetUnitInfo(unit): c_info = c_nvmlUnitInfo_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetUnitInfo") @@ -2544,13 +2632,15 @@ def nvmlUnitGetUnitInfo(unit): _nvmlCheckReturn(ret) return c_info + def nvmlUnitGetLedState(unit): - c_state = c_nvmlLedState_t() + c_state = c_nvmlLedState_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetLedState") ret = fn(unit, byref(c_state)) _nvmlCheckReturn(ret) return c_state + def nvmlUnitGetPsuInfo(unit): c_info = c_nvmlPSUInfo_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetPsuInfo") @@ -2558,6 +2648,7 @@ def nvmlUnitGetPsuInfo(unit): _nvmlCheckReturn(ret) return c_info + def nvmlUnitGetTemperature(unit, type): c_temp = c_uint() fn = _nvmlGetFunctionPointer("nvmlUnitGetTemperature") @@ -2565,6 +2656,7 @@ def nvmlUnitGetTemperature(unit, type): _nvmlCheckReturn(ret) return c_temp.value + def nvmlUnitGetFanSpeedInfo(unit): c_speeds = c_nvmlUnitFanSpeeds_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetFanSpeedInfo") @@ -2572,6 +2664,7 @@ def nvmlUnitGetFanSpeedInfo(unit): _nvmlCheckReturn(ret) return c_speeds + # added to API def nvmlUnitGetDeviceCount(unit): c_count = c_uint(0) @@ -2583,6 +2676,7 @@ def nvmlUnitGetDeviceCount(unit): _nvmlCheckReturn(ret) return c_count.value + def nvmlUnitGetDevices(unit): c_count = c_uint(nvmlUnitGetDeviceCount(unit)) device_array = c_nvmlDevice_t * c_count.value @@ -2592,6 +2686,7 @@ def nvmlUnitGetDevices(unit): _nvmlCheckReturn(ret) return c_devices + ## Device get functions def nvmlDeviceGetCount(): c_count = c_uint() @@ -2600,6 +2695,7 @@ def nvmlDeviceGetCount(): _nvmlCheckReturn(ret) return c_count.value + def nvmlDeviceGetHandleByIndex(index): c_index = c_uint(index) device = c_nvmlDevice_t() @@ -2608,6 +2704,7 @@ def nvmlDeviceGetHandleByIndex(index): _nvmlCheckReturn(ret) return device + @convertStrBytes def nvmlDeviceGetHandleBySerial(serial): c_serial = c_char_p(serial) @@ -2617,6 +2714,7 @@ def nvmlDeviceGetHandleBySerial(serial): _nvmlCheckReturn(ret) return device + @convertStrBytes def nvmlDeviceGetHandleByUUID(uuid): c_uuid = c_char_p(uuid) @@ -2626,6 +2724,7 @@ def nvmlDeviceGetHandleByUUID(uuid): _nvmlCheckReturn(ret) return device + @convertStrBytes def nvmlDeviceGetHandleByPciBusId(pciBusId): c_busId = c_char_p(pciBusId) @@ -2635,6 +2734,7 @@ def nvmlDeviceGetHandleByPciBusId(pciBusId): _nvmlCheckReturn(ret) return device + @convertStrBytes def nvmlDeviceGetName(handle): c_name = create_string_buffer(NVML_DEVICE_NAME_V2_BUFFER_SIZE) @@ -2643,14 +2743,17 @@ def nvmlDeviceGetName(handle): _nvmlCheckReturn(ret) return c_name.value + class c_nvmlDevicePerfModes_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), ('str', c_char * NVML_PERF_MODES_BUFFER_SIZE), ] + nvmlDevicePerfModes_v1 = 0x1000804 + @convertStrBytes def nvmlDeviceGetPerformanceModes(handle): perfModes = c_nvmlDevicePerfModes_v1_t() @@ -2660,14 +2763,17 @@ def nvmlDeviceGetPerformanceModes(handle): _nvmlCheckReturn(ret) return perfModes.str + class c_nvmlDeviceCurrentClockFreqs_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), ('str', c_char * NVML_PERF_MODES_BUFFER_SIZE), ] + nvmlDeviceCurrentClockFreqs_v1 = 0x1000804 + @convertStrBytes def nvmlDeviceGetCurrentClockFreqs(handle): currentClockFreqs = c_nvmlDeviceCurrentClockFreqs_v1_t() @@ -2677,20 +2783,23 @@ def nvmlDeviceGetCurrentClockFreqs(handle): _nvmlCheckReturn(ret) return currentClockFreqs.str + def nvmlDeviceGetBoardId(handle): - c_id = c_uint(); + c_id = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardId") ret = fn(handle, byref(c_id)) _nvmlCheckReturn(ret) return c_id.value + def nvmlDeviceGetMultiGpuBoard(handle): - c_multiGpu = c_uint(); + c_multiGpu = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMultiGpuBoard") ret = fn(handle, byref(c_multiGpu)) _nvmlCheckReturn(ret) return c_multiGpu.value + def nvmlDeviceGetBrand(handle): c_type = _nvmlBrandType_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBrand") @@ -2698,6 +2807,7 @@ def nvmlDeviceGetBrand(handle): _nvmlCheckReturn(ret) return c_type.value + def nvmlDeviceGetC2cModeInfoV1(handle): c_info = c_nvmlC2cModeInfo_v1_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetC2cModeInfoV") @@ -2705,9 +2815,11 @@ def nvmlDeviceGetC2cModeInfoV1(handle): _nvmlCheckReturn(ret) return c_info + def nvmlDeviceGetC2cModeInfoV(handle): return nvmlDeviceGetC2cModeInfoV1(handle) + @convertStrBytes def nvmlDeviceGetBoardPartNumber(handle): c_part_number = create_string_buffer(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE) @@ -2716,6 +2828,7 @@ def nvmlDeviceGetBoardPartNumber(handle): _nvmlCheckReturn(ret) return c_part_number.value + @convertStrBytes def nvmlDeviceGetSerial(handle): c_serial = create_string_buffer(NVML_DEVICE_SERIAL_BUFFER_SIZE) @@ -2724,6 +2837,7 @@ def nvmlDeviceGetSerial(handle): _nvmlCheckReturn(ret) return c_serial.value + def nvmlDeviceGetModuleId(handle, moduleId=c_uint()): isReference = type(moduleId) is not c_uint moduleIdRef = moduleId if isReference else byref(moduleId) @@ -2735,6 +2849,7 @@ def nvmlDeviceGetModuleId(handle, moduleId=c_uint()): _nvmlCheckReturn(ret) return moduleId.value + def nvmlDeviceGetMemoryAffinity(handle, nodeSetSize, scope): affinity_array = c_ulonglong * nodeSetSize c_affinity = affinity_array() @@ -2743,6 +2858,7 @@ def nvmlDeviceGetMemoryAffinity(handle, nodeSetSize, scope): _nvmlCheckReturn(ret) return c_affinity + def nvmlDeviceGetCpuAffinityWithinScope(handle, cpuSetSize, scope): affinity_array = c_ulonglong * cpuSetSize c_affinity = affinity_array() @@ -2751,6 +2867,7 @@ def nvmlDeviceGetCpuAffinityWithinScope(handle, cpuSetSize, scope): _nvmlCheckReturn(ret) return c_affinity + def nvmlDeviceGetCpuAffinity(handle, cpuSetSize): affinity_array = c_ulonglong * cpuSetSize c_affinity = affinity_array() @@ -2759,18 +2876,21 @@ def nvmlDeviceGetCpuAffinity(handle, cpuSetSize): _nvmlCheckReturn(ret) return c_affinity + def nvmlDeviceSetCpuAffinity(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceSetCpuAffinity") ret = fn(handle) _nvmlCheckReturn(ret) return None + def nvmlDeviceClearCpuAffinity(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceClearCpuAffinity") ret = fn(handle) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetNumaNodeId(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumaNodeId") node = c_int() @@ -2778,6 +2898,7 @@ def nvmlDeviceGetNumaNodeId(handle): _nvmlCheckReturn(ret) return node.value + def nvmlDeviceGetMinorNumber(handle): c_minor_number = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinorNumber") @@ -2785,6 +2906,7 @@ def nvmlDeviceGetMinorNumber(handle): _nvmlCheckReturn(ret) return c_minor_number.value + @convertStrBytes def nvmlDeviceGetUUID(handle): c_uuid = create_string_buffer(NVML_DEVICE_UUID_V2_BUFFER_SIZE) @@ -2793,15 +2915,16 @@ def nvmlDeviceGetUUID(handle): _nvmlCheckReturn(ret) return c_uuid.value + @convertStrBytes def nvmlDeviceGetInforomVersion(handle, infoRomObject): c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomVersion") - ret = fn(handle, _nvmlInforomObject_t(infoRomObject), - c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)) + ret = fn(handle, _nvmlInforomObject_t(infoRomObject), c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_version.value + # Added in 4.304 @convertStrBytes def nvmlDeviceGetInforomImageVersion(handle): @@ -2811,6 +2934,7 @@ def nvmlDeviceGetInforomImageVersion(handle): _nvmlCheckReturn(ret) return c_version.value + # Added in 4.304 def nvmlDeviceGetInforomConfigurationChecksum(handle): c_checksum = c_uint() @@ -2819,6 +2943,7 @@ def nvmlDeviceGetInforomConfigurationChecksum(handle): _nvmlCheckReturn(ret) return c_checksum.value + # Added in 4.304 def nvmlDeviceValidateInforom(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceValidateInforom") @@ -2826,6 +2951,7 @@ def nvmlDeviceValidateInforom(handle): _nvmlCheckReturn(ret) return None + def nvmlDeviceGetLastBBXFlushTime(handle): c_timestamp = c_ulonglong() c_durationUs = c_ulong() @@ -2834,6 +2960,7 @@ def nvmlDeviceGetLastBBXFlushTime(handle): _nvmlCheckReturn(ret) return [c_timestamp.value, c_durationUs.value] + def nvmlDeviceGetDisplayMode(handle): c_mode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayMode") @@ -2841,6 +2968,7 @@ def nvmlDeviceGetDisplayMode(handle): _nvmlCheckReturn(ret) return c_mode.value + def nvmlDeviceGetDisplayActive(handle): c_mode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayActive") @@ -2856,12 +2984,14 @@ def nvmlDeviceGetPersistenceMode(handle): _nvmlCheckReturn(ret) return c_state.value + def nvmlDeviceGetPciInfoExt(handle, c_info): fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfoExt") ret = fn(handle, c_info) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetPciInfo_v3(handle): c_info = nvmlPciInfo_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfo_v3") @@ -2869,9 +2999,11 @@ def nvmlDeviceGetPciInfo_v3(handle): _nvmlCheckReturn(ret) return c_info + def nvmlDeviceGetPciInfo(handle): return nvmlDeviceGetPciInfo_v3(handle) + def nvmlDeviceGetClockInfo(handle, type): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockInfo") @@ -2879,6 +3011,7 @@ def nvmlDeviceGetClockInfo(handle, type): _nvmlCheckReturn(ret) return c_clock.value + # Added in 2.285 def nvmlDeviceGetMaxClockInfo(handle, type): c_clock = c_uint() @@ -2887,6 +3020,7 @@ def nvmlDeviceGetMaxClockInfo(handle, type): _nvmlCheckReturn(ret) return c_clock.value + # Added in 4.304 def nvmlDeviceGetApplicationsClock(handle, type): c_clock = c_uint() @@ -2895,6 +3029,7 @@ def nvmlDeviceGetApplicationsClock(handle, type): _nvmlCheckReturn(ret) return c_clock.value + def nvmlDeviceGetMaxCustomerBoostClock(handle, type): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxCustomerBoostClock") @@ -2902,6 +3037,7 @@ def nvmlDeviceGetMaxCustomerBoostClock(handle, type): _nvmlCheckReturn(ret) return c_clock.value + def nvmlDeviceGetClock(handle, type, id): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetClock") @@ -2909,6 +3045,7 @@ def nvmlDeviceGetClock(handle, type, id): _nvmlCheckReturn(ret) return c_clock.value + # Added in 5.319 def nvmlDeviceGetDefaultApplicationsClock(handle, type): c_clock = c_uint() @@ -2917,6 +3054,7 @@ def nvmlDeviceGetDefaultApplicationsClock(handle, type): _nvmlCheckReturn(ret) return c_clock.value + # Added in 4.304 def nvmlDeviceGetSupportedMemoryClocks(handle): # first call to get the size @@ -2945,6 +3083,7 @@ def nvmlDeviceGetSupportedMemoryClocks(handle): # error case raise NVMLError(ret) + # Added in 4.304 def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): # first call to get the size @@ -2973,6 +3112,7 @@ def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): # error case raise NVMLError(ret) + def nvmlDeviceGetFanSpeed(handle): c_speed = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed") @@ -2980,6 +3120,7 @@ def nvmlDeviceGetFanSpeed(handle): _nvmlCheckReturn(ret) return c_speed.value + def nvmlDeviceGetFanSpeed_v2(handle, fan): c_speed = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed_v2") @@ -2987,6 +3128,7 @@ def nvmlDeviceGetFanSpeed_v2(handle, fan): _nvmlCheckReturn(ret) return c_speed.value + class c_nvmlFanSpeedInfo_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -2994,8 +3136,10 @@ class c_nvmlFanSpeedInfo_t(_PrintableStructure): ('speed', c_uint), ] + nvmlFanSpeedInfo_v1 = 0x100000C + def nvmlDeviceGetFanSpeedRPM(handle): c_fanSpeed = c_nvmlFanSpeedInfo_t() c_fanSpeed.fan = 0 @@ -3005,6 +3149,7 @@ def nvmlDeviceGetFanSpeedRPM(handle): _nvmlCheckReturn(ret) return c_fanSpeed.speed + def nvmlDeviceGetTargetFanSpeed(handle, fan): c_speed = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTargetFanSpeed") @@ -3012,6 +3157,7 @@ def nvmlDeviceGetTargetFanSpeed(handle, fan): _nvmlCheckReturn(ret) return c_speed.value + def nvmlDeviceGetNumFans(device): c_numFans = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumFans") @@ -3019,12 +3165,14 @@ def nvmlDeviceGetNumFans(device): _nvmlCheckReturn(ret) return c_numFans.value + def nvmlDeviceSetDefaultFanSpeed_v2(handle, index): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultFanSpeed_v2"); + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultFanSpeed_v2") ret = fn(handle, index) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetMinMaxFanSpeed(handle, minSpeed=c_uint(), maxSpeed=c_uint()): isReference = (type(minSpeed) is not c_uint) or (type(maxSpeed) is not c_uint) minSpeedRef = minSpeed if isReference else byref(minSpeed) @@ -3034,6 +3182,7 @@ def nvmlDeviceGetMinMaxFanSpeed(handle, minSpeed=c_uint(), maxSpeed=c_uint()): _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else [minSpeed.value, maxSpeed.value] + def nvmlDeviceGetFanControlPolicy_v2(handle, fan, fanControlPolicy=c_uint()): isReference = type(fanControlPolicy) is not c_uint fanControlPolicyRef = fanControlPolicy if isReference else byref(fanControlPolicy) @@ -3042,35 +3191,42 @@ def nvmlDeviceGetFanControlPolicy_v2(handle, fan, fanControlPolicy=c_uint()): _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else fanControlPolicy.value + def nvmlDeviceSetFanControlPolicy(handle, fan, fanControlPolicy): fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanControlPolicy") ret = fn(handle, fan, _nvmlFanControlPolicy_t(fanControlPolicy)) _nvmlCheckReturn(ret) return NVML_SUCCESS + class c_nvmlTemperature_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), ('sensorType', _nvmlTemperatureSensors_t), ('temperature', c_int), ] + + nvmlTemperature_v1 = 0x100000C + def nvmlDeviceGetTemperatureV1(handle, sensor): c_temp = c_nvmlTemperature_v1_t() c_temp.version = nvmlTemperature_v1 - c_temp.sensorType = _nvmlTemperatureSensors_t(sensor) + c_temp.sensorType = _nvmlTemperatureSensors_t(sensor) fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureV") ret = fn(handle, byref(c_temp)) _nvmlCheckReturn(ret) return c_temp.temperature + def nvmlDeviceGetTemperatureV(handle, sensor, version=nvmlTemperature_v1): if version == nvmlTemperature_v1: return nvmlDeviceGetTemperatureV1(handle, sensor) else: raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + # DEPRECATED use nvmlDeviceGetTemperatureV instead def nvmlDeviceGetTemperature(handle, sensor): c_temp = c_uint() @@ -3079,6 +3235,7 @@ def nvmlDeviceGetTemperature(handle, sensor): _nvmlCheckReturn(ret) return c_temp.value + def nvmlDeviceGetTemperatureThreshold(handle, threshold): c_temp = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureThreshold") @@ -3086,6 +3243,7 @@ def nvmlDeviceGetTemperatureThreshold(handle, threshold): _nvmlCheckReturn(ret) return c_temp.value + def nvmlDeviceSetTemperatureThreshold(handle, threshold, temp): c_temp = c_uint() c_temp.value = temp @@ -3094,6 +3252,7 @@ def nvmlDeviceSetTemperatureThreshold(handle, threshold, temp): _nvmlCheckReturn(ret) return None + def nvmlDeviceGetMarginTemperature(handle): c_marginTempInfo = c_nvmlMarginTemperature_v1_t() c_marginTempInfo.version = nvmlMarginTemperature_v1 @@ -3102,6 +3261,7 @@ def nvmlDeviceGetMarginTemperature(handle): _nvmlCheckReturn(ret) return c_marginTempInfo.marginTemperature + # DEPRECATED use nvmlDeviceGetPerformanceState def nvmlDeviceGetPowerState(handle): c_pstate = _nvmlPstates_t() @@ -3110,6 +3270,7 @@ def nvmlDeviceGetPowerState(handle): _nvmlCheckReturn(ret) return c_pstate.value + def nvmlDeviceGetPerformanceState(handle): c_pstate = _nvmlPstates_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceState") @@ -3117,6 +3278,7 @@ def nvmlDeviceGetPerformanceState(handle): _nvmlCheckReturn(ret) return c_pstate.value + def nvmlDeviceGetPowerManagementMode(handle): c_pcapMode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementMode") @@ -3124,6 +3286,7 @@ def nvmlDeviceGetPowerManagementMode(handle): _nvmlCheckReturn(ret) return c_pcapMode.value + def nvmlDeviceGetPowerManagementLimit(handle): c_limit = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimit") @@ -3131,6 +3294,7 @@ def nvmlDeviceGetPowerManagementLimit(handle): _nvmlCheckReturn(ret) return c_limit.value + # Added in 4.304 def nvmlDeviceGetPowerManagementLimitConstraints(handle): c_minLimit = c_uint() @@ -3140,6 +3304,7 @@ def nvmlDeviceGetPowerManagementLimitConstraints(handle): _nvmlCheckReturn(ret) return [c_minLimit.value, c_maxLimit.value] + # Added in 4.304 def nvmlDeviceGetPowerManagementDefaultLimit(handle): c_limit = c_uint() @@ -3157,6 +3322,7 @@ def nvmlDeviceGetEnforcedPowerLimit(handle): _nvmlCheckReturn(ret) return c_limit.value + def nvmlDeviceGetPowerUsage(handle): c_watts = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerUsage") @@ -3164,6 +3330,7 @@ def nvmlDeviceGetPowerUsage(handle): _nvmlCheckReturn(ret) return c_watts.value + def nvmlDeviceGetTotalEnergyConsumption(handle): c_millijoules = c_uint64() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEnergyConsumption") @@ -3171,6 +3338,7 @@ def nvmlDeviceGetTotalEnergyConsumption(handle): _nvmlCheckReturn(ret) return c_millijoules.value + # Added in 4.304 def nvmlDeviceGetGpuOperationMode(handle): c_currState = _nvmlGpuOperationMode_t() @@ -3180,14 +3348,17 @@ def nvmlDeviceGetGpuOperationMode(handle): _nvmlCheckReturn(ret) return [c_currState.value, c_pendingState.value] + # Added in 4.304 def nvmlDeviceGetCurrentGpuOperationMode(handle): return nvmlDeviceGetGpuOperationMode(handle)[0] + # Added in 4.304 def nvmlDeviceGetPendingGpuOperationMode(handle): return nvmlDeviceGetGpuOperationMode(handle)[1] + def nvmlDeviceGetMemoryInfo(handle, version=None): if not version: c_memory = c_nvmlMemory_t() @@ -3200,6 +3371,7 @@ def nvmlDeviceGetMemoryInfo(handle, version=None): _nvmlCheckReturn(ret) return c_memory + def nvmlDeviceGetBAR1MemoryInfo(handle): c_bar1_memory = c_nvmlBAR1Memory_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBAR1MemoryInfo") @@ -3207,6 +3379,7 @@ def nvmlDeviceGetBAR1MemoryInfo(handle): _nvmlCheckReturn(ret) return c_bar1_memory + def nvmlDeviceGetComputeMode(handle): c_mode = _nvmlComputeMode_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeMode") @@ -3214,6 +3387,7 @@ def nvmlDeviceGetComputeMode(handle): _nvmlCheckReturn(ret) return c_mode.value + def nvmlDeviceGetCudaComputeCapability(handle): c_major = c_int() c_minor = c_int() @@ -3222,6 +3396,7 @@ def nvmlDeviceGetCudaComputeCapability(handle): _nvmlCheckReturn(ret) return (c_major.value, c_minor.value) + def nvmlDeviceGetEccMode(handle): c_currState = _nvmlEnableState_t() c_pendingState = _nvmlEnableState_t() @@ -3230,14 +3405,17 @@ def nvmlDeviceGetEccMode(handle): _nvmlCheckReturn(ret) return [c_currState.value, c_pendingState.value] + # added to API def nvmlDeviceGetCurrentEccMode(handle): return nvmlDeviceGetEccMode(handle)[0] + # added to API def nvmlDeviceGetPendingEccMode(handle): return nvmlDeviceGetEccMode(handle)[1] + def nvmlDeviceGetDefaultEccMode(handle): c_defaultState = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultEccMode") @@ -3245,35 +3423,34 @@ def nvmlDeviceGetDefaultEccMode(handle): _nvmlCheckReturn(ret) return [c_defaultState.value] + def nvmlDeviceGetTotalEccErrors(handle, errorType, counterType): c_count = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEccErrors") - ret = fn(handle, _nvmlMemoryErrorType_t(errorType), - _nvmlEccCounterType_t(counterType), byref(c_count)) + ret = fn(handle, _nvmlMemoryErrorType_t(errorType), _nvmlEccCounterType_t(counterType), byref(c_count)) _nvmlCheckReturn(ret) return c_count.value + # This is deprecated, instead use nvmlDeviceGetMemoryErrorCounter def nvmlDeviceGetDetailedEccErrors(handle, errorType, counterType): c_counts = c_nvmlEccErrorCounts_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDetailedEccErrors") - ret = fn(handle, _nvmlMemoryErrorType_t(errorType), - _nvmlEccCounterType_t(counterType), byref(c_counts)) + ret = fn(handle, _nvmlMemoryErrorType_t(errorType), _nvmlEccCounterType_t(counterType), byref(c_counts)) _nvmlCheckReturn(ret) return c_counts + # Added in 4.304 def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType): c_count = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryErrorCounter") - ret = fn(handle, - _nvmlMemoryErrorType_t(errorType), - _nvmlEccCounterType_t(counterType), - _nvmlMemoryLocation_t(locationType), - byref(c_count)) + ret = fn(handle, _nvmlMemoryErrorType_t(errorType), _nvmlEccCounterType_t(counterType), + _nvmlMemoryLocation_t(locationType), byref(c_count)) _nvmlCheckReturn(ret) return c_count.value + def nvmlDeviceGetUtilizationRates(handle): c_util = c_nvmlUtilization_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetUtilizationRates") @@ -3281,6 +3458,7 @@ def nvmlDeviceGetUtilizationRates(handle): _nvmlCheckReturn(ret) return c_util + def nvmlDeviceGetEncoderUtilization(handle): c_util = c_uint() c_samplingPeriod = c_uint() @@ -3289,6 +3467,7 @@ def nvmlDeviceGetEncoderUtilization(handle): _nvmlCheckReturn(ret) return [c_util.value, c_samplingPeriod.value] + def nvmlDeviceGetDecoderUtilization(handle): c_util = c_uint() c_samplingPeriod = c_uint() @@ -3297,6 +3476,7 @@ def nvmlDeviceGetDecoderUtilization(handle): _nvmlCheckReturn(ret) return [c_util.value, c_samplingPeriod.value] + def nvmlDeviceGetJpgUtilization(handle): c_util = c_uint() c_samplingPeriod = c_uint() @@ -3305,6 +3485,7 @@ def nvmlDeviceGetJpgUtilization(handle): _nvmlCheckReturn(ret) return [c_util.value, c_samplingPeriod.value] + def nvmlDeviceGetOfaUtilization(handle): c_util = c_uint() c_samplingPeriod = c_uint() @@ -3313,6 +3494,7 @@ def nvmlDeviceGetOfaUtilization(handle): _nvmlCheckReturn(ret) return [c_util.value, c_samplingPeriod.value] + def nvmlDeviceGetPcieReplayCounter(handle): c_replay = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieReplayCounter") @@ -3320,6 +3502,7 @@ def nvmlDeviceGetPcieReplayCounter(handle): _nvmlCheckReturn(ret) return c_replay.value + def nvmlDeviceGetDriverModel(handle): c_currModel = _nvmlDriverModel_t() c_pendingModel = _nvmlDriverModel_t() @@ -3328,14 +3511,17 @@ def nvmlDeviceGetDriverModel(handle): _nvmlCheckReturn(ret) return [c_currModel.value, c_pendingModel.value] + # added to API def nvmlDeviceGetCurrentDriverModel(handle): return nvmlDeviceGetDriverModel(handle)[0] + # added to API def nvmlDeviceGetPendingDriverModel(handle): return nvmlDeviceGetDriverModel(handle)[1] + # Added in 2.285 @convertStrBytes def nvmlDeviceGetVbiosVersion(handle): @@ -3345,6 +3531,7 @@ def nvmlDeviceGetVbiosVersion(handle): _nvmlCheckReturn(ret) return c_version.value + # Added in 2.285 def nvmlDeviceGetComputeRunningProcesses_v2(handle): # first call to get the size @@ -3376,6 +3563,7 @@ def nvmlDeviceGetComputeRunningProcesses_v2(handle): # error case raise NVMLError(ret) + # Added in 2.285 def nvmlDeviceGetComputeRunningProcesses_v3(handle): # first call to get the size @@ -3411,10 +3599,12 @@ def nvmlDeviceGetComputeRunningProcesses_v3(handle): # error case raise NVMLError(ret) + @throwOnVersionMismatch def nvmlDeviceGetComputeRunningProcesses(handle): return nvmlDeviceGetComputeRunningProcesses_v3(handle) + def nvmlDeviceGetGraphicsRunningProcesses_v2(handle): # first call to get the size c_count = c_uint(0) @@ -3445,6 +3635,7 @@ def nvmlDeviceGetGraphicsRunningProcesses_v2(handle): # error case raise NVMLError(ret) + def nvmlDeviceGetGraphicsRunningProcesses_v3(handle): # first call to get the size c_count = c_uint(0) @@ -3479,14 +3670,17 @@ def nvmlDeviceGetGraphicsRunningProcesses_v3(handle): # error case raise NVMLError(ret) + @throwOnVersionMismatch def nvmlDeviceGetGraphicsRunningProcesses(handle): return nvmlDeviceGetGraphicsRunningProcesses_v3(handle) + @throwOnVersionMismatch def nvmlDeviceGetMPSComputeRunningProcesses(handle): return nvmlDeviceGetMPSComputeRunningProcesses_v3(handle) + def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle): # first call to get the size c_count = c_uint(0) @@ -3521,6 +3715,7 @@ def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle): # error case raise NVMLError(ret) + def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle): # first call to get the size c_count = c_uint(0) @@ -3555,6 +3750,7 @@ def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle): # error case raise NVMLError(ret) + def nvmlDeviceGetRunningProcessDetailList(handle, version, mode): c_processDetailList = c_nvmlProcessDetailList_t() c_processDetailList.version = version @@ -3590,6 +3786,7 @@ def nvmlDeviceGetRunningProcessDetailList(handle, version, mode): # error case raise NVMLError(ret) + def nvmlDeviceGetAutoBoostedClocksEnabled(handle): c_isEnabled = _nvmlEnableState_t() c_defaultIsEnabled = _nvmlEnableState_t() @@ -3599,6 +3796,7 @@ def nvmlDeviceGetAutoBoostedClocksEnabled(handle): return [c_isEnabled.value, c_defaultIsEnabled.value] #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + ## Set functions def nvmlUnitSetLedState(unit, color): fn = _nvmlGetFunctionPointer("nvmlUnitSetLedState") @@ -3606,36 +3804,42 @@ def nvmlUnitSetLedState(unit, color): _nvmlCheckReturn(ret) return None + def nvmlDeviceSetPersistenceMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetPersistenceMode") ret = fn(handle, _nvmlEnableState_t(mode)) _nvmlCheckReturn(ret) return None + def nvmlDeviceSetComputeMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetComputeMode") ret = fn(handle, _nvmlComputeMode_t(mode)) _nvmlCheckReturn(ret) return None + def nvmlDeviceSetEccMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetEccMode") ret = fn(handle, _nvmlEnableState_t(mode)) _nvmlCheckReturn(ret) return None + def nvmlDeviceClearEccErrorCounts(handle, counterType): fn = _nvmlGetFunctionPointer("nvmlDeviceClearEccErrorCounts") ret = fn(handle, _nvmlEccCounterType_t(counterType)) _nvmlCheckReturn(ret) return None + def nvmlDeviceSetDriverModel(handle, model): fn = _nvmlGetFunctionPointer("nvmlDeviceSetDriverModel") ret = fn(handle, _nvmlDriverModel_t(model)) _nvmlCheckReturn(ret) return None + def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled): fn = _nvmlGetFunctionPointer("nvmlDeviceSetAutoBoostedClocksEnabled") ret = fn(handle, _nvmlEnableState_t(enabled)) @@ -3643,6 +3847,7 @@ def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled): return None #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags): fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultAutoBoostedClocksEnabled") ret = fn(handle, _nvmlEnableState_t(enabled), c_uint(flags)) @@ -3650,30 +3855,35 @@ def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags): return None #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + def nvmlDeviceSetGpuLockedClocks(handle, minGpuClockMHz, maxGpuClockMHz): fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuLockedClocks") ret = fn(handle, c_uint(minGpuClockMHz), c_uint(maxGpuClockMHz)) _nvmlCheckReturn(ret) return None + def nvmlDeviceResetGpuLockedClocks(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceResetGpuLockedClocks") ret = fn(handle) _nvmlCheckReturn(ret) return None + def nvmlDeviceSetMemoryLockedClocks(handle, minMemClockMHz, maxMemClockMHz): fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemoryLockedClocks") ret = fn(handle, c_uint(minMemClockMHz), c_uint(maxMemClockMHz)) _nvmlCheckReturn(ret) return None + def nvmlDeviceResetMemoryLockedClocks(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceResetMemoryLockedClocks") ret = fn(handle) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetClkMonStatus(handle, c_clkMonInfo=nvmlClkMonStatus_t()): isReference = type(c_clkMonInfo) is not nvmlClkMonStatus_t c_clkMonInfoRef = c_clkMonInfo if isReference else byref(c_clkMonInfo) @@ -3682,6 +3892,7 @@ def nvmlDeviceGetClkMonStatus(handle, c_clkMonInfo=nvmlClkMonStatus_t()): _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else c_clkMonInfo + # Added in 4.304 def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz): fn = _nvmlGetFunctionPointer("nvmlDeviceSetApplicationsClocks") @@ -3689,6 +3900,7 @@ def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz) _nvmlCheckReturn(ret) return None + # Added in 4.304 def nvmlDeviceResetApplicationsClocks(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceResetApplicationsClocks") @@ -3696,6 +3908,7 @@ def nvmlDeviceResetApplicationsClocks(handle): _nvmlCheckReturn(ret) return None + # Added in 4.304 def nvmlDeviceSetPowerManagementLimit(handle, limit): fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit") @@ -3703,6 +3916,7 @@ def nvmlDeviceSetPowerManagementLimit(handle, limit): _nvmlCheckReturn(ret) return None + # Added in 4.304 def nvmlDeviceSetGpuOperationMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuOperationMode") @@ -3710,6 +3924,7 @@ def nvmlDeviceSetGpuOperationMode(handle, mode): _nvmlCheckReturn(ret) return None + # Added in 2.285 def nvmlEventSetCreate(): fn = _nvmlGetFunctionPointer("nvmlEventSetCreate") @@ -3718,6 +3933,7 @@ def nvmlEventSetCreate(): _nvmlCheckReturn(ret) return eventSet + # Added in 2.285 def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet): fn = _nvmlGetFunctionPointer("nvmlDeviceRegisterEvents") @@ -3725,6 +3941,7 @@ def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet): _nvmlCheckReturn(ret) return None + # Added in 2.285 def nvmlDeviceGetSupportedEventTypes(handle): c_eventTypes = c_ulonglong() @@ -3733,6 +3950,7 @@ def nvmlDeviceGetSupportedEventTypes(handle): _nvmlCheckReturn(ret) return c_eventTypes.value + # raises NVML_ERROR_TIMEOUT exception on timeout def nvmlEventSetWait_v2(eventSet, timeoutms): fn = _nvmlGetFunctionPointer("nvmlEventSetWait_v2") @@ -3741,9 +3959,11 @@ def nvmlEventSetWait_v2(eventSet, timeoutms): _nvmlCheckReturn(ret) return data + def nvmlEventSetWait(eventSet, timeoutms): return nvmlEventSetWait_v2(eventSet, timeoutms) + # Added in 2.285 def nvmlEventSetFree(eventSet): fn = _nvmlGetFunctionPointer("nvmlEventSetFree") @@ -3751,6 +3971,7 @@ def nvmlEventSetFree(eventSet): _nvmlCheckReturn(ret) return None + # Added in 3.295 def nvmlDeviceOnSameBoard(handle1, handle2): fn = _nvmlGetFunctionPointer("nvmlDeviceOnSameBoard") @@ -3759,6 +3980,7 @@ def nvmlDeviceOnSameBoard(handle1, handle2): _nvmlCheckReturn(ret) return (onSameBoard.value != 0) + # Added in 3.295 def nvmlDeviceGetCurrPcieLinkGeneration(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkGeneration") @@ -3767,6 +3989,7 @@ def nvmlDeviceGetCurrPcieLinkGeneration(handle): _nvmlCheckReturn(ret) return gen.value + # Added in 3.295 def nvmlDeviceGetMaxPcieLinkGeneration(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkGeneration") @@ -3775,6 +3998,7 @@ def nvmlDeviceGetMaxPcieLinkGeneration(handle): _nvmlCheckReturn(ret) return gen.value + # Added in 3.295 def nvmlDeviceGetCurrPcieLinkWidth(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkWidth") @@ -3783,6 +4007,7 @@ def nvmlDeviceGetCurrPcieLinkWidth(handle): _nvmlCheckReturn(ret) return width.value + # Added in 3.295 def nvmlDeviceGetMaxPcieLinkWidth(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkWidth") @@ -3791,6 +4016,7 @@ def nvmlDeviceGetMaxPcieLinkWidth(handle): _nvmlCheckReturn(ret) return width.value + def nvmlDeviceGetGpuMaxPcieLinkGeneration(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuMaxPcieLinkGeneration") gen = c_uint() @@ -3798,36 +4024,41 @@ def nvmlDeviceGetGpuMaxPcieLinkGeneration(handle): _nvmlCheckReturn(ret) return gen.value + # Added in 4.304 def nvmlDeviceGetSupportedClocksThrottleReasons(handle): - c_reasons= c_ulonglong() + c_reasons = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksThrottleReasons") ret = fn(handle, byref(c_reasons)) _nvmlCheckReturn(ret) return c_reasons.value + def nvmlDeviceGetSupportedClocksEventReasons(handle): - c_reasons= c_ulonglong() + c_reasons = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksEventReasons") ret = fn(handle, byref(c_reasons)) _nvmlCheckReturn(ret) return c_reasons.value + # Added in 4.304 def nvmlDeviceGetCurrentClocksThrottleReasons(handle): - c_reasons= c_ulonglong() + c_reasons = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksThrottleReasons") ret = fn(handle, byref(c_reasons)) _nvmlCheckReturn(ret) return c_reasons.value + def nvmlDeviceGetCurrentClocksEventReasons(handle): - c_reasons= c_ulonglong() + c_reasons = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksEventReasons") ret = fn(handle, byref(c_reasons)) _nvmlCheckReturn(ret) return c_reasons.value + # Added in 5.319 def nvmlDeviceGetIndex(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetIndex") @@ -3836,6 +4067,7 @@ def nvmlDeviceGetIndex(handle): _nvmlCheckReturn(ret) return c_index.value + # Added in 5.319 def nvmlDeviceGetAccountingMode(handle): c_mode = _nvmlEnableState_t() @@ -3844,18 +4076,21 @@ def nvmlDeviceGetAccountingMode(handle): _nvmlCheckReturn(ret) return c_mode.value + def nvmlDeviceSetAccountingMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetAccountingMode") ret = fn(handle, _nvmlEnableState_t(mode)) _nvmlCheckReturn(ret) return None + def nvmlDeviceClearAccountingPids(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceClearAccountingPids") ret = fn(handle) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetAccountingStats(handle, pid): stats = c_nvmlAccountingStats_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingStats") @@ -3866,6 +4101,7 @@ def nvmlDeviceGetAccountingStats(handle, pid): stats.maxMemoryUsage = None return stats + def nvmlDeviceGetAccountingPids(handle): count = c_uint(nvmlDeviceGetAccountingBufferSize(handle)) pids = (c_uint * count.value)() @@ -3874,6 +4110,7 @@ def nvmlDeviceGetAccountingPids(handle): _nvmlCheckReturn(ret) return list(map(int, pids[0:count.value])) + def nvmlDeviceGetAccountingBufferSize(handle): bufferSize = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingBufferSize") @@ -3881,6 +4118,7 @@ def nvmlDeviceGetAccountingBufferSize(handle): _nvmlCheckReturn(ret) return int(bufferSize.value) + def nvmlDeviceGetRetiredPages(device, sourceFilter): c_source = _nvmlPageRetirementCause_t(sourceFilter) c_count = c_uint(0) @@ -3890,8 +4128,7 @@ def nvmlDeviceGetRetiredPages(device, sourceFilter): ret = fn(device, c_source, byref(c_count), None) # this should only fail with insufficient size - if ((ret != NVML_SUCCESS) and - (ret != NVML_ERROR_INSUFFICIENT_SIZE)): + if ((ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE)): raise NVMLError(ret) # call again with a buffer @@ -3904,6 +4141,7 @@ def nvmlDeviceGetRetiredPages(device, sourceFilter): _nvmlCheckReturn(ret) return list(map(int, c_pages[0:c_count.value])) + def nvmlDeviceGetRetiredPages_v2(device, sourceFilter): c_source = _nvmlPageRetirementCause_t(sourceFilter) c_count = c_uint(0) @@ -3913,8 +4151,7 @@ def nvmlDeviceGetRetiredPages_v2(device, sourceFilter): ret = fn(device, c_source, byref(c_count), None) # this should only fail with insufficient size - if ((ret != NVML_SUCCESS) and - (ret != NVML_ERROR_INSUFFICIENT_SIZE)): + if ((ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE)): raise NVMLError(ret) # call again with a buffer @@ -3927,7 +4164,8 @@ def nvmlDeviceGetRetiredPages_v2(device, sourceFilter): c_times = times_array() ret = fn(device, c_source, byref(c_count), c_pages, c_times) _nvmlCheckReturn(ret) - return [ { 'address': int(c_pages[i]), 'timestamp': int(c_times[i]) } for i in range(c_count.value) ]; + return [{'address': int(c_pages[i]), 'timestamp': int(c_times[i])} for i in range(c_count.value)] + def nvmlDeviceGetRetiredPagesPendingStatus(device): c_pending = _nvmlEnableState_t() @@ -3936,6 +4174,7 @@ def nvmlDeviceGetRetiredPagesPendingStatus(device): _nvmlCheckReturn(ret) return int(c_pending.value) + def nvmlDeviceGetAPIRestriction(device, apiType): c_permission = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAPIRestriction") @@ -3943,12 +4182,14 @@ def nvmlDeviceGetAPIRestriction(device, apiType): _nvmlCheckReturn(ret) return int(c_permission.value) + def nvmlDeviceSetAPIRestriction(handle, apiType, isRestricted): fn = _nvmlGetFunctionPointer("nvmlDeviceSetAPIRestriction") ret = fn(handle, _nvmlRestrictedAPI_t(apiType), _nvmlEnableState_t(isRestricted)) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetBridgeChipInfo(handle): bridgeHierarchy = c_nvmlBridgeChipHierarchy_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBridgeChipInfo") @@ -3956,6 +4197,7 @@ def nvmlDeviceGetBridgeChipInfo(handle): _nvmlCheckReturn(ret) return bridgeHierarchy + def nvmlDeviceGetSamples(device, sampling_type, timeStamp): c_sampling_type = _nvmlSamplingType_t(sampling_type) c_time_stamp = c_ulonglong(timeStamp) @@ -3972,10 +4214,11 @@ def nvmlDeviceGetSamples(device, sampling_type, timeStamp): sampleArray = c_sample_count.value * c_nvmlSample_t c_samples = sampleArray() - ret = fn(device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), c_samples) + ret = fn(device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), c_samples) _nvmlCheckReturn(ret) return (c_sample_value_type.value, c_samples[0:c_sample_count.value]) + def nvmlDeviceGetViolationStatus(device, perfPolicyType): c_perfPolicy_type = _nvmlPerfPolicyType_t(perfPolicyType) c_violTime = c_nvmlViolationTime_t() @@ -3986,6 +4229,7 @@ def nvmlDeviceGetViolationStatus(device, perfPolicyType): _nvmlCheckReturn(ret) return c_violTime + def nvmlDeviceGetPcieThroughput(device, counter): c_util = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieThroughput") @@ -3993,6 +4237,7 @@ def nvmlDeviceGetPcieThroughput(device, counter): _nvmlCheckReturn(ret) return c_util.value + def nvmlSystemGetTopologyGpuSet(cpuNumber): c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlSystemGetTopologyGpuSet") @@ -4009,6 +4254,7 @@ def nvmlSystemGetTopologyGpuSet(cpuNumber): _nvmlCheckReturn(ret) return list(c_devices[0:c_count.value]) + def nvmlDeviceGetTopologyNearestGpus(device, level): c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyNearestGpus") @@ -4026,6 +4272,7 @@ def nvmlDeviceGetTopologyNearestGpus(device, level): _nvmlCheckReturn(ret) return list(c_devices[0:c_count.value]) + def nvmlDeviceGetTopologyCommonAncestor(device1, device2): c_level = _nvmlGpuTopologyLevel_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyCommonAncestor") @@ -4033,6 +4280,7 @@ def nvmlDeviceGetTopologyCommonAncestor(device1, device2): _nvmlCheckReturn(ret) return c_level.value + def nvmlDeviceGetNvLinkUtilizationCounter(device, link, counter): c_rxcounter = c_ulonglong() c_txcounter = c_ulonglong() @@ -4041,24 +4289,28 @@ def nvmlDeviceGetNvLinkUtilizationCounter(device, link, counter): _nvmlCheckReturn(ret) return (c_rxcounter.value, c_txcounter.value) + def nvmlDeviceFreezeNvLinkUtilizationCounter(device, link, counter, freeze): fn = _nvmlGetFunctionPointer("nvmlDeviceFreezeNvLinkUtilizationCounter") ret = fn(device, link, counter, freeze) _nvmlCheckReturn(ret) return None + def nvmlDeviceResetNvLinkUtilizationCounter(device, link, counter): fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkUtilizationCounter") ret = fn(device, link, counter) _nvmlCheckReturn(ret) return None + def nvmlDeviceSetNvLinkUtilizationControl(device, link, counter, control, reset): fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkUtilizationControl") ret = fn(device, link, counter, byref(control), reset) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetNvLinkUtilizationControl(device, link, counter): c_control = nvmlNvLinkUtilizationControl_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationControl") @@ -4066,6 +4318,7 @@ def nvmlDeviceGetNvLinkUtilizationControl(device, link, counter): _nvmlCheckReturn(ret) return c_control + def nvmlDeviceGetNvLinkCapability(device, link, capability): c_capResult = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkCapability") @@ -4073,6 +4326,7 @@ def nvmlDeviceGetNvLinkCapability(device, link, capability): _nvmlCheckReturn(ret) return c_capResult.value + def nvmlDeviceGetNvLinkErrorCounter(device, link, counter): c_result = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkErrorCounter") @@ -4080,12 +4334,14 @@ def nvmlDeviceGetNvLinkErrorCounter(device, link, counter): _nvmlCheckReturn(ret) return c_result.value + def nvmlDeviceResetNvLinkErrorCounters(device, link): fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkErrorCounters") ret = fn(device, link) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetNvLinkRemotePciInfo(device, link): c_pci = nvmlPciInfo_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemotePciInfo_v2") @@ -4093,6 +4349,7 @@ def nvmlDeviceGetNvLinkRemotePciInfo(device, link): _nvmlCheckReturn(ret) return c_pci + def nvmlDeviceGetNvLinkRemoteDeviceType(handle, link): c_type = _nvmlNvLinkDeviceType_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemoteDeviceType") @@ -4100,6 +4357,7 @@ def nvmlDeviceGetNvLinkRemoteDeviceType(handle, link): _nvmlCheckReturn(ret) return c_type.value + def nvmlDeviceGetNvLinkState(device, link): c_isActive = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkState") @@ -4107,6 +4365,7 @@ def nvmlDeviceGetNvLinkState(device, link): _nvmlCheckReturn(ret) return c_isActive.value + def nvmlDeviceGetNvLinkVersion(device, link): c_version = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkVersion") @@ -4114,12 +4373,14 @@ def nvmlDeviceGetNvLinkVersion(device, link): _nvmlCheckReturn(ret) return c_version.value + def nvmlDeviceModifyDrainState(pciInfo, newState): fn = _nvmlGetFunctionPointer("nvmlDeviceModifyDrainState") ret = fn(pointer(pciInfo), newState) _nvmlCheckReturn(ret) return None + def nvmlDeviceQueryDrainState(pciInfo): c_newState = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceQueryDrainState") @@ -4127,18 +4388,21 @@ def nvmlDeviceQueryDrainState(pciInfo): _nvmlCheckReturn(ret) return c_newState.value + def nvmlDeviceRemoveGpu(pciInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceRemoveGpu") ret = fn(pointer(pciInfo)) _nvmlCheckReturn(ret) return None + def nvmlDeviceDiscoverGpus(pciInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceDiscoverGpus") ret = fn(pointer(pciInfo)) _nvmlCheckReturn(ret) return None + def nvmlDeviceGetFieldValues(handle, fieldIds): values_arr = c_nvmlFieldValue_t * len(fieldIds) values = values_arr() @@ -4154,6 +4418,7 @@ def nvmlDeviceGetFieldValues(handle, fieldIds): _nvmlCheckReturn(ret) return values + def nvmlDeviceClearFieldValues(handle, fieldIds): values_arr = c_nvmlFieldValue_t * len(fieldIds) values = values_arr() @@ -4169,6 +4434,7 @@ def nvmlDeviceClearFieldValues(handle, fieldIds): _nvmlCheckReturn(ret) return values + def nvmlDeviceGetVirtualizationMode(handle): c_virtualization_mode = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVirtualizationMode") @@ -4176,10 +4442,12 @@ def nvmlDeviceGetVirtualizationMode(handle): _nvmlCheckReturn(ret) return c_virtualization_mode.value + def nvmlDeviceSetVirtualizationMode(handle, virtualization_mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetVirtualizationMode") return fn(handle, virtualization_mode) + def nvmlDeviceGetVgpuHeterogeneousMode(handle): c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 @@ -4188,6 +4456,7 @@ def nvmlDeviceGetVgpuHeterogeneousMode(handle): _nvmlCheckReturn(ret) return c_vgpuHeterogeneousMode.mode + def nvmlDeviceSetVgpuHeterogeneousMode(handle, heterogeneous_mode): c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 @@ -4197,6 +4466,7 @@ def nvmlDeviceSetVgpuHeterogeneousMode(handle, heterogeneous_mode): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlVgpuInstanceGetPlacementId(vgpuInstance): c_placement = c_nvmlVgpuPlacementId_v1_t(0) c_placement.version = VgpuPlacementId_v1 @@ -4205,9 +4475,10 @@ def nvmlVgpuInstanceGetPlacementId(vgpuInstance): _nvmlCheckReturn(ret) return c_placement.placementId + def nvmlDeviceGetVgpuTypeSupportedPlacements(handle, vgpuTypeId, mode=0, version=1): c_max_instances = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") ret = fn(handle, vgpuTypeId, byref(c_max_instances)) _nvmlCheckReturn(ret) @@ -4224,14 +4495,15 @@ def nvmlDeviceGetVgpuTypeSupportedPlacements(handle, vgpuTypeId, mode=0, version c_placements = c_uint * c_max_instances.value c_vgpu_placements.placementIds = c_placements() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeSupportedPlacements") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeSupportedPlacements") ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) _nvmlCheckReturn(ret) return c_vgpu_placements + def nvmlDeviceGetVgpuTypeCreatablePlacements(handle, vgpuTypeId, version=1): c_max_instances = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") ret = fn(handle, vgpuTypeId, byref(c_max_instances)) _nvmlCheckReturn(ret) @@ -4245,11 +4517,12 @@ def nvmlDeviceGetVgpuTypeCreatablePlacements(handle, vgpuTypeId, version=1): c_placements = c_uint * c_max_instances.value c_vgpu_placements.placementIds = c_placements() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeCreatablePlacements") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeCreatablePlacements") ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) _nvmlCheckReturn(ret) return c_vgpu_placements + def nvmlGetVgpuDriverCapabilities(capability): c_capResult = c_uint() fn = _nvmlGetFunctionPointer("nvmlGetVgpuDriverCapabilities") @@ -4257,6 +4530,7 @@ def nvmlGetVgpuDriverCapabilities(capability): _nvmlCheckReturn(ret) return c_capResult.value + def nvmlDeviceGetVgpuCapabilities(handle, capability): c_capResult = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuCapabilities") @@ -4264,17 +4538,19 @@ def nvmlDeviceGetVgpuCapabilities(handle, capability): _nvmlCheckReturn(ret) return c_capResult.value + def nvmlDeviceSetVgpuCapabilities(handle, capability, state): fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuCapabilities") ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), state) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetSupportedVgpus(handle): # first call to get the size c_vgpu_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedVgpus") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedVgpus") ret = fn(handle, byref(c_vgpu_count), None) if (ret == NVML_SUCCESS): @@ -4296,11 +4572,12 @@ def nvmlDeviceGetSupportedVgpus(handle): # error case raise NVMLError(ret) + def nvmlDeviceGetCreatableVgpus(handle): # first call to get the size c_vgpu_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetCreatableVgpus") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCreatableVgpus") ret = fn(handle, byref(c_vgpu_count), None) if (ret == NVML_SUCCESS): @@ -4322,126 +4599,142 @@ def nvmlDeviceGetCreatableVgpus(handle): # error case raise NVMLError(ret) + def nvmlVgpuTypeGetGpuInstanceProfileId(vgpuTypeId): c_profile_id = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGpuInstanceProfileId") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGpuInstanceProfileId") ret = fn(vgpuTypeId, byref(c_profile_id)) _nvmlCheckReturn(ret) return (c_profile_id.value) + @convertStrBytes def nvmlVgpuTypeGetClass(vgpuTypeId): c_class = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetClass") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetClass") ret = fn(vgpuTypeId, c_class, byref(c_buffer_size)) _nvmlCheckReturn(ret) return c_class.value + @convertStrBytes def nvmlVgpuTypeGetName(vgpuTypeId): c_name = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetName") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetName") ret = fn(vgpuTypeId, c_name, byref(c_buffer_size)) _nvmlCheckReturn(ret) return c_name.value + def nvmlVgpuTypeGetDeviceID(vgpuTypeId): - c_device_id = c_ulonglong(0) + c_device_id = c_ulonglong(0) c_subsystem_id = c_ulonglong(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetDeviceID") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetDeviceID") ret = fn(vgpuTypeId, byref(c_device_id), byref(c_subsystem_id)) _nvmlCheckReturn(ret) return (c_device_id.value, c_subsystem_id.value) + def nvmlVgpuTypeGetFramebufferSize(vgpuTypeId): c_fb_size = c_ulonglong(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFramebufferSize") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFramebufferSize") ret = fn(vgpuTypeId, byref(c_fb_size)) _nvmlCheckReturn(ret) return c_fb_size.value + def nvmlVgpuTypeGetNumDisplayHeads(vgpuTypeId): c_num_heads = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetNumDisplayHeads") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetNumDisplayHeads") ret = fn(vgpuTypeId, byref(c_num_heads)) _nvmlCheckReturn(ret) return c_num_heads.value + def nvmlVgpuTypeGetResolution(vgpuTypeId): c_xdim = c_uint(0) c_ydim = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetResolution") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetResolution") ret = fn(vgpuTypeId, 0, byref(c_xdim), byref(c_ydim)) _nvmlCheckReturn(ret) return (c_xdim.value, c_ydim.value) + @convertStrBytes def nvmlVgpuTypeGetLicense(vgpuTypeId): c_license = create_string_buffer(NVML_GRID_LICENSE_BUFFER_SIZE) c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetLicense") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetLicense") ret = fn(vgpuTypeId, c_license, c_buffer_size) _nvmlCheckReturn(ret) return c_license.value + def nvmlVgpuTypeGetFrameRateLimit(vgpuTypeId): c_frl_config = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFrameRateLimit") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFrameRateLimit") ret = fn(vgpuTypeId, byref(c_frl_config)) _nvmlCheckReturn(ret) return c_frl_config.value + def nvmlVgpuTypeGetGspHeapSize(vgpuTypeId): c_gsp_heap = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGspHeapSize") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGspHeapSize") ret = fn(vgpuTypeId, byref(c_gsp_heap)) _nvmlCheckReturn(ret) return c_gsp_heap.value + def nvmlVgpuTypeGetFbReservation(vgpuTypeId): c_fb_reservation = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFbReservation") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFbReservation") ret = fn(vgpuTypeId, byref(c_fb_reservation)) _nvmlCheckReturn(ret) return c_fb_reservation.value + def nvmlVgpuInstanceGetRuntimeStateSize(vgpuInstance): c_runtime_state = nvmlVgpuRuntimeState_v1_t() c_runtime_state.version = VgpuRuntimeState_v1 - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetRuntimeStateSize") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetRuntimeStateSize") ret = fn(vgpuInstance, byref(c_runtime_state)) _nvmlCheckReturn(ret) return c_runtime_state + def nvmlVgpuTypeGetMaxInstances(handle, vgpuTypeId): c_max_instances = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") ret = fn(handle, vgpuTypeId, byref(c_max_instances)) _nvmlCheckReturn(ret) return c_max_instances.value + def nvmlVgpuTypeGetMaxInstancesPerVm(vgpuTypeId): c_max_instances_per_vm = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerVm") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerVm") ret = fn(vgpuTypeId, byref(c_max_instances_per_vm)) _nvmlCheckReturn(ret) return c_max_instances_per_vm.value + def nvmlVgpuTypeGetBAR1Info(vgpuTypeId): c_bar1Info = c_nvmlVgpuTypeBar1Info_v1_t(0) c_bar1Info.version = VgpuTypeBar1Info_v1 - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetBAR1Info") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetBAR1Info") ret = fn(vgpuTypeId, byref(c_bar1Info)) _nvmlCheckReturn(ret) return c_bar1Info + def nvmlDeviceGetActiveVgpus(handle): # first call to get the size c_vgpu_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetActiveVgpus") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetActiveVgpus") ret = fn(handle, byref(c_vgpu_count), None) if (ret == NVML_SUCCESS): @@ -4463,67 +4756,76 @@ def nvmlDeviceGetActiveVgpus(handle): # error case raise NVMLError(ret) + @convertStrBytes def nvmlVgpuInstanceGetVmID(vgpuInstance): c_vm_id = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) - c_vm_id_type = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmID") + c_vm_id_type = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmID") ret = fn(vgpuInstance, byref(c_vm_id), c_buffer_size, byref(c_vm_id_type)) _nvmlCheckReturn(ret) return (c_vm_id.value, c_vm_id_type.value) + @convertStrBytes def nvmlVgpuInstanceGetUUID(vgpuInstance): c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetUUID") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetUUID") ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) _nvmlCheckReturn(ret) return c_uuid.value + @convertStrBytes def nvmlVgpuInstanceGetMdevUUID(vgpuInstance): c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMdevUUID") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMdevUUID") ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) _nvmlCheckReturn(ret) return c_uuid.value + @convertStrBytes def nvmlVgpuInstanceGetVmDriverVersion(vgpuInstance): c_driver_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) c_buffer_size = c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmDriverVersion") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmDriverVersion") ret = fn(vgpuInstance, byref(c_driver_version), c_buffer_size) _nvmlCheckReturn(ret) return c_driver_version.value + def nvmlVgpuInstanceGetLicenseStatus(vgpuInstance): c_license_status = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseStatus") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseStatus") ret = fn(vgpuInstance, byref(c_license_status)) _nvmlCheckReturn(ret) return c_license_status.value + def nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance): - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseInfo_v2") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseInfo_v2") c_license_info = c_nvmlVgpuLicenseInfo_t() ret = fn(vgpuInstance, byref(c_license_info)) _nvmlCheckReturn(ret) return c_license_info + def nvmlVgpuInstanceGetLicenseInfo(vgpuInstance): return nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance) + def nvmlVgpuInstanceGetFrameRateLimit(vgpuInstance): c_frl = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFrameRateLimit") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFrameRateLimit") ret = fn(vgpuInstance, byref(c_frl)) _nvmlCheckReturn(ret) return c_frl.value + def nvmlVgpuInstanceGetEccMode(vgpuInstance): c_mode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEccMode") @@ -4531,45 +4833,52 @@ def nvmlVgpuInstanceGetEccMode(vgpuInstance): _nvmlCheckReturn(ret) return c_mode.value + def nvmlVgpuInstanceGetType(vgpuInstance): c_vgpu_type = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetType") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetType") ret = fn(vgpuInstance, byref(c_vgpu_type)) _nvmlCheckReturn(ret) return c_vgpu_type.value + def nvmlVgpuInstanceGetEncoderCapacity(vgpuInstance): c_encoder_capacity = c_ulonglong(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderCapacity") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderCapacity") ret = fn(vgpuInstance, byref(c_encoder_capacity)) _nvmlCheckReturn(ret) return c_encoder_capacity.value + def nvmlVgpuInstanceSetEncoderCapacity(vgpuInstance, encoder_capacity): - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceSetEncoderCapacity") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceSetEncoderCapacity") return fn(vgpuInstance, encoder_capacity) + def nvmlVgpuInstanceGetFbUsage(vgpuInstance): c_fb_usage = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFbUsage") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFbUsage") ret = fn(vgpuInstance, byref(c_fb_usage)) _nvmlCheckReturn(ret) return c_fb_usage.value + def nvmlVgpuTypeGetCapabilities(vgpuTypeId, capability): c_cap_result = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetCapabilities") + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetCapabilities") ret = fn(vgpuTypeId, _nvmlVgpuCapability_t(capability), byref(c_cap_result)) _nvmlCheckReturn(ret) return (c_cap_result.value) + def nvmlVgpuInstanceGetGpuInstanceId(vgpuInstance): c_id = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuInstanceId") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuInstanceId") ret = fn(vgpuInstance, byref(c_id)) _nvmlCheckReturn(ret) return (c_id.value) + @convertStrBytes def nvmlVgpuInstanceGetGpuPciId(vgpuInstance): c_vgpuPciId = create_string_buffer(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE) @@ -4578,13 +4887,14 @@ def nvmlVgpuInstanceGetGpuPciId(vgpuInstance): _nvmlCheckReturn(ret) return c_vgpuPciId.value + def nvmlDeviceGetVgpuUtilization(handle, timeStamp): # first call to get the size c_vgpu_count = c_uint(0) c_time_stamp = c_ulonglong(timeStamp) c_sample_value_type = _nvmlValueType_t() - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuUtilization") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuUtilization") ret = fn(handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), None) if (ret == NVML_SUCCESS): @@ -4604,6 +4914,7 @@ def nvmlDeviceGetVgpuUtilization(handle, timeStamp): # error case raise NVMLError(ret) + def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp): # first call to get the size c_time_stamp = c_ulonglong(timeStamp) @@ -4613,7 +4924,7 @@ def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp): c_vgpuUtilInfo.vgpuInstanceCount = c_uint(0) c_vgpuUtilInfo.lastSeenTimeStamp = c_time_stamp - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuInstancesUtilizationInfo") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuInstancesUtilizationInfo") ret = fn(handle, byref(c_vgpuUtilInfo)) if (ret == NVML_SUCCESS): @@ -4634,13 +4945,15 @@ def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp): # error case raise NVMLError(ret) + def nvmlDeviceGetP2PStatus(device1, device2, p2pIndex): c_p2pstatus = _nvmlGpuP2PStatus_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetP2PStatus") - ret = fn(device1, device2,p2pIndex, byref(c_p2pstatus)) + ret = fn(device1, device2, p2pIndex, byref(c_p2pstatus)) _nvmlCheckReturn(ret) return c_p2pstatus.value + def nvmlDeviceGetGridLicensableFeatures_v4(handle): c_get_grid_licensable_features = c_nvmlGridLicensableFeatures_v4_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGridLicensableFeatures_v4") @@ -4649,9 +4962,11 @@ def nvmlDeviceGetGridLicensableFeatures_v4(handle): return (c_get_grid_licensable_features) + def nvmlDeviceGetGridLicensableFeatures(handle): return nvmlDeviceGetGridLicensableFeatures_v4(handle) + def nvmlDeviceGetGspFirmwareVersion(handle, version=None): isUserDefined = version is not None if not isUserDefined: @@ -4661,6 +4976,7 @@ def nvmlDeviceGetGspFirmwareVersion(handle, version=None): _nvmlCheckReturn(ret) return NVML_SUCCESS if isUserDefined else version.value + def nvmlDeviceGetGspFirmwareMode(handle, isEnabled=c_uint(), defaultMode=c_uint()): isReference = type(isEnabled) is not c_uint isEnabledRef = isEnabled if isReference else byref(isEnabled) @@ -4670,6 +4986,7 @@ def nvmlDeviceGetGspFirmwareMode(handle, isEnabled=c_uint(), defaultMode=c_uint( _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else [isEnabled.value, defaultMode.value] + def nvmlDeviceGetEncoderCapacity(handle, encoderQueryType): c_encoder_capacity = c_ulonglong(0) c_encoderQuery_type = _nvmlEncoderQueryType_t(encoderQueryType) @@ -4679,12 +4996,13 @@ def nvmlDeviceGetEncoderCapacity(handle, encoderQueryType): _nvmlCheckReturn(ret) return c_encoder_capacity.value + def nvmlDeviceGetVgpuProcessUtilization(handle, timeStamp): # first call to get the size c_vgpu_count = c_uint(0) c_time_stamp = c_ulonglong(timeStamp) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessUtilization") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessUtilization") ret = fn(handle, c_time_stamp, byref(c_vgpu_count), None) if (ret == NVML_SUCCESS): @@ -4704,6 +5022,7 @@ def nvmlDeviceGetVgpuProcessUtilization(handle, timeStamp): # error case raise NVMLError(ret) + def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp): # first call to get the size c_time_stamp = c_ulonglong(timeStamp) @@ -4712,7 +5031,7 @@ def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp): c_vgpuProcUtilInfo.vgpuProcessCount = c_uint(0) c_vgpuProcUtilInfo.lastSeenTimeStamp = c_time_stamp - fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessesUtilizationInfo") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessesUtilizationInfo") ret = fn(handle, byref(c_vgpuProcUtilInfo)) if (ret == NVML_SUCCESS): @@ -4733,6 +5052,7 @@ def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp): # error case raise NVMLError(ret) + def nvmlDeviceGetEncoderStats(handle): c_encoderCount = c_ulonglong(0) c_encodeFps = c_ulonglong(0) @@ -4742,11 +5062,12 @@ def nvmlDeviceGetEncoderStats(handle): _nvmlCheckReturn(ret) return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) + def nvmlDeviceGetEncoderSessions(handle): # first call to get the size c_session_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderSessions") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderSessions") ret = fn(handle, byref(c_session_count), None) if (ret == NVML_SUCCESS): @@ -4768,6 +5089,7 @@ def nvmlDeviceGetEncoderSessions(handle): # error case raise NVMLError(ret) + def nvmlDeviceGetFBCStats(handle): c_fbcStats = c_nvmlFBCStats_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCStats") @@ -4775,11 +5097,12 @@ def nvmlDeviceGetFBCStats(handle): _nvmlCheckReturn(ret) return c_fbcStats + def nvmlDeviceGetFBCSessions(handle): # first call to get the size c_session_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCSessions") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCSessions") ret = fn(handle, byref(c_session_count), None) if (ret == NVML_SUCCESS): @@ -4801,20 +5124,22 @@ def nvmlDeviceGetFBCSessions(handle): # error case raise NVMLError(ret) + def nvmlVgpuInstanceGetEncoderStats(vgpuInstance): - c_encoderCount = c_ulonglong(0) - c_encodeFps = c_ulonglong(0) - c_encoderLatency = c_ulonglong(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderStats") + c_encoderCount = c_ulonglong(0) + c_encodeFps = c_ulonglong(0) + c_encoderLatency = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderStats") ret = fn(vgpuInstance, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency)) _nvmlCheckReturn(ret) return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) + def nvmlVgpuInstanceGetEncoderSessions(vgpuInstance): # first call to get the size c_session_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderSessions") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderSessions") ret = fn(vgpuInstance, byref(c_session_count), None) if (ret == NVML_SUCCESS): @@ -4836,6 +5161,7 @@ def nvmlVgpuInstanceGetEncoderSessions(vgpuInstance): # error case raise NVMLError(ret) + def nvmlVgpuInstanceGetFBCStats(vgpuInstance): c_fbcStats = c_nvmlFBCStats_t() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCStats") @@ -4843,11 +5169,12 @@ def nvmlVgpuInstanceGetFBCStats(vgpuInstance): _nvmlCheckReturn(ret) return c_fbcStats + def nvmlVgpuInstanceGetFBCSessions(vgpuInstance): # first call to get the size c_session_count = c_uint(0) - fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCSessions") + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCSessions") ret = fn(vgpuInstance, byref(c_session_count), None) if (ret == NVML_SUCCESS): @@ -4869,12 +5196,13 @@ def nvmlVgpuInstanceGetFBCSessions(vgpuInstance): # error case raise NVMLError(ret) + def nvmlDeviceGetProcessUtilization(handle, timeStamp): # first call to get the size c_count = c_uint(0) c_time_stamp = c_ulonglong(timeStamp) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessUtilization") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessUtilization") ret = fn(handle, None, byref(c_count), c_time_stamp) if (ret == NVML_ERROR_INSUFFICIENT_SIZE): @@ -4891,6 +5219,7 @@ def nvmlDeviceGetProcessUtilization(handle, timeStamp): # error case raise NVMLError(ret) + def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp): # first call to get the size c_time_stamp = c_ulonglong(timeStamp) @@ -4899,7 +5228,7 @@ def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp): c_processesUtilInfo.processSamplesCount = c_uint(0) c_processesUtilInfo.lastSeenTimeStamp = c_time_stamp - fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessesUtilizationInfo") + fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessesUtilizationInfo") ret = fn(handle, byref(c_processesUtilInfo)) if (ret == NVML_ERROR_INSUFFICIENT_SIZE): @@ -4917,6 +5246,7 @@ def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp): # error case raise NVMLError(ret) + def nvmlVgpuInstanceGetMetadata(vgpuInstance): fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMetadata") c_vgpuMetadata = c_nvmlVgpuMetadata_t() @@ -4931,6 +5261,7 @@ def nvmlVgpuInstanceGetMetadata(vgpuInstance): raise NVMLError(ret) return c_vgpuMetadata + def nvmlDeviceGetVgpuMetadata(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuMetadata") c_vgpuPgpuMetadata = c_nvmlVgpuPgpuMetadata_t() @@ -4945,6 +5276,7 @@ def nvmlDeviceGetVgpuMetadata(handle): raise NVMLError(ret) return c_vgpuPgpuMetadata + def nvmlGetVgpuCompatibility(vgpuMetadata, pgpuMetadata): fn = _nvmlGetFunctionPointer("nvmlGetVgpuCompatibility") c_vgpuPgpuCompatibility = c_nvmlVgpuPgpuCompatibility_t() @@ -4952,6 +5284,7 @@ def nvmlGetVgpuCompatibility(vgpuMetadata, pgpuMetadata): _nvmlCheckReturn(ret) return c_vgpuPgpuCompatibility + @convertStrBytes def nvmlDeviceGetPgpuMetadataString(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetPgpuMetadataString") @@ -4967,6 +5300,7 @@ def nvmlDeviceGetPgpuMetadataString(handle): raise NVMLError(ret) return (c_pgpuMetadata.value, c_bufferSize.value) + def nvmlDeviceGetVgpuSchedulerLog(handle): c_vgpu_sched_log = c_nvmlVgpuSchedulerLog_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerLog") @@ -4974,6 +5308,7 @@ def nvmlDeviceGetVgpuSchedulerLog(handle): _nvmlCheckReturn(ret) return c_vgpu_sched_log + def nvmlDeviceGetVgpuSchedulerState(handle): c_vgpu_sched_state = c_nvmlVgpuSchedulerGetState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerState") @@ -4981,6 +5316,7 @@ def nvmlDeviceGetVgpuSchedulerState(handle): _nvmlCheckReturn(ret) return c_vgpu_sched_state + def nvmlDeviceGetVgpuSchedulerCapabilities(handle): c_vgpu_sched_caps = c_nvmlVgpuSchedulerCapabilities_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerCapabilities") @@ -4988,18 +5324,21 @@ def nvmlDeviceGetVgpuSchedulerCapabilities(handle): _nvmlCheckReturn(ret) return c_vgpu_sched_caps + def nvmlDeviceSetVgpuSchedulerState(handle, sched_state): fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuSchedulerState") ret = fn(handle, byref(sched_state)) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlSetVgpuVersion(vgpuVersion): fn = _nvmlGetFunctionPointer("nvmlSetVgpuVersion") ret = fn(byref(vgpuVersion)) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlGetVgpuVersion(supported=None, current=None): isUserDefined = (supported is not None) or (current is not None) if not isUserDefined: @@ -5008,10 +5347,9 @@ def nvmlGetVgpuVersion(supported=None, current=None): fn = _nvmlGetFunctionPointer("nvmlGetVgpuVersion") ret = fn(byref(supported), byref(current)) _nvmlCheckReturn(ret) - return NVML_SUCCESS if isUserDefined else [(supported.minVersion, - supported.maxVersion), - (current.minVersion, - current.maxVersion)] + return NVML_SUCCESS if isUserDefined else [(supported.minVersion, supported.maxVersion), + (current.minVersion, current.maxVersion)] + def nvmlVgpuInstanceGetAccountingMode(vgpuInstance): c_mode = _nvmlEnableState_t() @@ -5020,6 +5358,7 @@ def nvmlVgpuInstanceGetAccountingMode(vgpuInstance): _nvmlCheckReturn(ret) return c_mode.value + def nvmlVgpuInstanceGetAccountingPids(vgpuInstance): c_pidCount = c_uint() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingPids") @@ -5033,6 +5372,7 @@ def nvmlVgpuInstanceGetAccountingPids(vgpuInstance): raise NVMLError(ret) return (c_pidCount, c_pidArray) + def nvmlVgpuInstanceGetAccountingStats(vgpuInstance, pid): c_accountingStats = c_nvmlAccountingStats_t() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingStats") @@ -5040,12 +5380,14 @@ def nvmlVgpuInstanceGetAccountingStats(vgpuInstance, pid): _nvmlCheckReturn(ret) return c_accountingStats + def nvmlVgpuInstanceClearAccountingPids(vgpuInstance): fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceClearAccountingPids") ret = fn(vgpuInstance) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlGetExcludedDeviceCount(): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceCount") @@ -5053,6 +5395,7 @@ def nvmlGetExcludedDeviceCount(): _nvmlCheckReturn(ret) return c_count.value + def nvmlGetExcludedDeviceInfoByIndex(index): c_index = c_uint(index) info = c_nvmlExcludedDeviceInfo_t() @@ -5061,6 +5404,7 @@ def nvmlGetExcludedDeviceInfoByIndex(index): _nvmlCheckReturn(ret) return info + def nvmlDeviceGetHostVgpuMode(handle): c_host_vgpu_mode = _nvmlHostVgpuMode_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetHostVgpuMode") @@ -5068,6 +5412,7 @@ def nvmlDeviceGetHostVgpuMode(handle): _nvmlCheckReturn(ret) return c_host_vgpu_mode.value + def nvmlDeviceSetMigMode(device, mode): c_activationStatus = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceSetMigMode") @@ -5075,6 +5420,7 @@ def nvmlDeviceSetMigMode(device, mode): _nvmlCheckReturn(ret) return c_activationStatus.value + def nvmlDeviceGetMigMode(device): c_currentMode = c_uint() c_pendingMode = c_uint() @@ -5083,6 +5429,7 @@ def nvmlDeviceGetMigMode(device): _nvmlCheckReturn(ret) return [c_currentMode.value, c_pendingMode.value] + def nvmlDeviceGetGpuInstanceProfileInfo(device, profile, version=2): if version == 2: c_info = c_nvmlGpuInstanceProfileInfo_v2_t() @@ -5096,9 +5443,11 @@ def nvmlDeviceGetGpuInstanceProfileInfo(device, profile, version=2): _nvmlCheckReturn(ret) return c_info + # Define function alias for the API exposed by NVML nvmlDeviceGetGpuInstanceProfileInfoV = nvmlDeviceGetGpuInstanceProfileInfo + def nvmlDeviceGetGpuInstanceRemainingCapacity(device, profileId): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceRemainingCapacity") @@ -5106,12 +5455,14 @@ def nvmlDeviceGetGpuInstanceRemainingCapacity(device, profileId): _nvmlCheckReturn(ret) return c_count.value + def nvmlDeviceGetGpuInstancePossiblePlacements(device, profileId, placementsRef, countRef): fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstancePossiblePlacements_v2") ret = fn(device, profileId, placementsRef, countRef) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceCreateGpuInstance(device, profileId): c_instance = c_nvmlGpuInstance_t() fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstance") @@ -5119,6 +5470,7 @@ def nvmlDeviceCreateGpuInstance(device, profileId): _nvmlCheckReturn(ret) return c_instance + def nvmlDeviceCreateGpuInstanceWithPlacement(device, profileId, placement): c_instance = c_nvmlGpuInstance_t() fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstanceWithPlacement") @@ -5126,18 +5478,21 @@ def nvmlDeviceCreateGpuInstanceWithPlacement(device, profileId, placement): _nvmlCheckReturn(ret) return c_instance + def nvmlGpuInstanceDestroy(gpuInstance): fn = _nvmlGetFunctionPointer("nvmlGpuInstanceDestroy") ret = fn(gpuInstance) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetGpuInstances(device, profileId, gpuInstancesRef, countRef): fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstances") ret = fn(device, profileId, gpuInstancesRef, countRef) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetGpuInstanceById(device, gpuInstanceId): c_instance = c_nvmlGpuInstance_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceById") @@ -5145,6 +5500,7 @@ def nvmlDeviceGetGpuInstanceById(device, gpuInstanceId): _nvmlCheckReturn(ret) return c_instance + def nvmlGpuInstanceGetInfo(gpuInstance): c_info = c_nvmlGpuInstanceInfo_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetInfo") @@ -5152,6 +5508,7 @@ def nvmlGpuInstanceGetInfo(gpuInstance): _nvmlCheckReturn(ret) return c_info + def nvmlGpuInstanceGetComputeInstanceProfileInfo(device, profile, engProfile, version=2): if version == 2: c_info = c_nvmlComputeInstanceProfileInfo_v2_t() @@ -5160,14 +5517,16 @@ def nvmlGpuInstanceGetComputeInstanceProfileInfo(device, profile, engProfile, ve c_info = c_nvmlComputeInstanceProfileInfo_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfo") else: - raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) ret = fn(device, profile, engProfile, byref(c_info)) _nvmlCheckReturn(ret) return c_info + # Define function alias for the API exposed by NVML nvmlGpuInstanceGetComputeInstanceProfileInfoV = nvmlGpuInstanceGetComputeInstanceProfileInfo + def nvmlGpuInstanceGetComputeInstanceRemainingCapacity(gpuInstance, profileId): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceRemainingCapacity") @@ -5175,12 +5534,14 @@ def nvmlGpuInstanceGetComputeInstanceRemainingCapacity(gpuInstance, profileId): _nvmlCheckReturn(ret) return c_count.value + def nvmlGpuInstanceGetComputeInstancePossiblePlacements(gpuInstance, profileId, placementsRef, countRef): fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstancePossiblePlacements") ret = fn(gpuInstance, profileId, placementsRef, countRef) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlGpuInstanceCreateComputeInstance(gpuInstance, profileId): c_instance = c_nvmlComputeInstance_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstance") @@ -5188,6 +5549,7 @@ def nvmlGpuInstanceCreateComputeInstance(gpuInstance, profileId): _nvmlCheckReturn(ret) return c_instance + def nvmlGpuInstanceCreateComputeInstanceWithPlacement(gpuInstance, profileId, placement): c_instance = c_nvmlComputeInstance_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstanceWithPlacement") @@ -5195,18 +5557,21 @@ def nvmlGpuInstanceCreateComputeInstanceWithPlacement(gpuInstance, profileId, pl _nvmlCheckReturn(ret) return c_instance + def nvmlComputeInstanceDestroy(computeInstance): fn = _nvmlGetFunctionPointer("nvmlComputeInstanceDestroy") ret = fn(computeInstance) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlGpuInstanceGetComputeInstances(gpuInstance, profileId, computeInstancesRef, countRef): fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstances") ret = fn(gpuInstance, profileId, computeInstancesRef, countRef) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlGpuInstanceGetComputeInstanceById(gpuInstance, computeInstanceId): c_instance = c_nvmlComputeInstance_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceById") @@ -5214,6 +5579,7 @@ def nvmlGpuInstanceGetComputeInstanceById(gpuInstance, computeInstanceId): _nvmlCheckReturn(ret) return c_instance + def nvmlComputeInstanceGetInfo_v2(computeInstance): c_info = c_nvmlComputeInstanceInfo_t() fn = _nvmlGetFunctionPointer("nvmlComputeInstanceGetInfo_v2") @@ -5221,9 +5587,11 @@ def nvmlComputeInstanceGetInfo_v2(computeInstance): _nvmlCheckReturn(ret) return c_info + def nvmlComputeInstanceGetInfo(computeInstance): return nvmlComputeInstanceGetInfo_v2(computeInstance) + def nvmlDeviceIsMigDeviceHandle(device): c_isMigDevice = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceIsMigDeviceHandle") @@ -5231,6 +5599,7 @@ def nvmlDeviceIsMigDeviceHandle(device): _nvmlCheckReturn(ret) return c_isMigDevice + def nvmlDeviceGetGpuInstanceId(device): c_gpuInstanceId = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceId") @@ -5238,6 +5607,7 @@ def nvmlDeviceGetGpuInstanceId(device): _nvmlCheckReturn(ret) return c_gpuInstanceId.value + def nvmlDeviceGetComputeInstanceId(device): c_computeInstanceId = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeInstanceId") @@ -5245,6 +5615,7 @@ def nvmlDeviceGetComputeInstanceId(device): _nvmlCheckReturn(ret) return c_computeInstanceId.value + def nvmlDeviceGetMaxMigDeviceCount(device): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxMigDeviceCount") @@ -5252,6 +5623,7 @@ def nvmlDeviceGetMaxMigDeviceCount(device): _nvmlCheckReturn(ret) return c_count.value + def nvmlDeviceGetMigDeviceHandleByIndex(device, index): c_index = c_uint(index) migDevice = c_nvmlDevice_t() @@ -5260,6 +5632,7 @@ def nvmlDeviceGetMigDeviceHandleByIndex(device, index): _nvmlCheckReturn(ret) return migDevice + def nvmlDeviceGetDeviceHandleFromMigDeviceHandle(migDevice): device = c_nvmlDevice_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDeviceHandleFromMigDeviceHandle") @@ -5267,6 +5640,7 @@ def nvmlDeviceGetDeviceHandleFromMigDeviceHandle(migDevice): _nvmlCheckReturn(ret) return device + def nvmlDeviceGetAttributes_v2(device): c_attrs = c_nvmlDeviceAttributes() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAttributes_v2") @@ -5274,9 +5648,11 @@ def nvmlDeviceGetAttributes_v2(device): _nvmlCheckReturn(ret) return c_attrs + def nvmlDeviceGetAttributes(device): return nvmlDeviceGetAttributes_v2(device) + def nvmlDeviceGetRemappedRows(device): fn = _nvmlGetFunctionPointer("nvmlDeviceGetRemappedRows") c_corr = c_uint() @@ -5287,6 +5663,7 @@ def nvmlDeviceGetRemappedRows(device): _nvmlCheckReturn(ret) return (c_corr.value, c_unc.value, c_bpending.value, c_bfailure.value) + def nvmlDeviceGetRowRemapperHistogram(device): c_vals = c_nvmlRowRemapperHistogramValues() fn = _nvmlGetFunctionPointer("nvmlDeviceGetRowRemapperHistogram") @@ -5294,6 +5671,7 @@ def nvmlDeviceGetRowRemapperHistogram(device): _nvmlCheckReturn(ret) return c_vals + def nvmlDeviceGetArchitecture(device): arch = _nvmlDeviceArchitecture_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetArchitecture") @@ -5301,6 +5679,7 @@ def nvmlDeviceGetArchitecture(device): _nvmlCheckReturn(ret) return arch.value + def nvmlDeviceGetBusType(device): c_busType = _nvmlBusType_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBusType") @@ -5308,6 +5687,7 @@ def nvmlDeviceGetBusType(device): _nvmlCheckReturn(ret) return c_busType.value + def nvmlDeviceGetIrqNum(device): c_irqNum = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetIrqNum") @@ -5315,6 +5695,7 @@ def nvmlDeviceGetIrqNum(device): _nvmlCheckReturn(ret) return c_irqNum.value + def nvmlDeviceGetNumGpuCores(device): c_numCores = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumGpuCores") @@ -5322,6 +5703,7 @@ def nvmlDeviceGetNumGpuCores(device): _nvmlCheckReturn(ret) return c_numCores.value + def nvmlDeviceGetPowerSource(device): c_powerSource = _nvmlPowerSource_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerSource") @@ -5329,6 +5711,7 @@ def nvmlDeviceGetPowerSource(device): _nvmlCheckReturn(ret) return c_powerSource.value + def nvmlDeviceGetMemoryBusWidth(device): c_memBusWidth = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryBusWidth") @@ -5336,6 +5719,7 @@ def nvmlDeviceGetMemoryBusWidth(device): _nvmlCheckReturn(ret) return c_memBusWidth.value + def nvmlDeviceGetPcieLinkMaxSpeed(device): c_speed = _nvmlPcieLinkMaxSpeed_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieLinkMaxSpeed") @@ -5343,6 +5727,7 @@ def nvmlDeviceGetPcieLinkMaxSpeed(device): _nvmlCheckReturn(ret) return c_speed.value + def nvmlDeviceGetAdaptiveClockInfoStatus(device): c_adaptiveClockInfoStatus = _nvmlAdaptiveClockInfoStatus_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAdaptiveClockInfoStatus") @@ -5350,6 +5735,7 @@ def nvmlDeviceGetAdaptiveClockInfoStatus(device): _nvmlCheckReturn(ret) return c_adaptiveClockInfoStatus.value + def nvmlDeviceGetPcieSpeed(device): c_speed = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieSpeed") @@ -5357,38 +5743,43 @@ def nvmlDeviceGetPcieSpeed(device): _nvmlCheckReturn(ret) return c_speed.value + def nvmlDeviceGetDynamicPstatesInfo(device, c_dynamicpstatesinfo=c_nvmlGpuDynamicPstatesInfo_t()): isReference = type(c_dynamicpstatesinfo) is not c_nvmlGpuDynamicPstatesInfo_t dynamicpstatesinfoRef = c_dynamicpstatesinfo if isReference else byref(c_dynamicpstatesinfo) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetDynamicPstatesInfo"); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDynamicPstatesInfo") ret = fn(device, dynamicpstatesinfoRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else c_dynamicpstatesinfo + def nvmlDeviceSetFanSpeed_v2(handle, index, speed): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanSpeed_v2"); + fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanSpeed_v2") ret = fn(handle, index, speed) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetThermalSettings(device, sensorindex, c_thermalsettings=c_nvmlGpuThermalSettings_t()): isReference = type(c_thermalsettings) is not c_nvmlGpuThermalSettings_t thermalsettingsRef = c_thermalsettings if isReference else byref(c_thermalsettings) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetThermalSettings"); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetThermalSettings") ret = fn(device, sensorindex, thermalsettingsRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else c_thermalsettings.sensor[:] + def nvmlDeviceGetMinMaxClockOfPState(device, clockType, pstate, minClockMHz=c_uint(), maxClockMHz=c_uint()): isReference = (type(minClockMHz) is not c_uint) or (type(maxClockMHz) is not c_uint) minClockMHzRef = minClockMHz if isReference else byref(minClockMHz) maxClockMHzRef = maxClockMHz if isReference else byref(maxClockMHz) - fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxClockOfPState"); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxClockOfPState") ret = fn(device, _nvmlClockType_t(clockType), _nvmlClockType_t(pstate), minClockMHzRef, maxClockMHzRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else (minClockMHz.value, maxClockMHz.value) + class c_nvmlClockOffset_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -5399,22 +5790,26 @@ class c_nvmlClockOffset_t(_PrintableStructure): ('maxClockOffsetMHz', c_int), ] + nvmlClockOffset_v1 = 0x1000018 + def nvmlDeviceGetClockOffsets(device, info): - fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockOffsets"); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockOffsets") ret = fn(device, info) return NVML_SUCCESS + def nvmlDeviceSetClockOffsets(device, info): - fn = _nvmlGetFunctionPointer("nvmlDeviceSetClockOffsets"); + fn = _nvmlGetFunctionPointer("nvmlDeviceSetClockOffsets") ret = fn(device, info) return NVML_SUCCESS + def nvmlDeviceGetSupportedPerformanceStates(device): pstates = [] c_count = c_uint(NVML_MAX_GPU_PERF_PSTATES) - c_size = sizeof(c_uint)*c_count.value + c_size = sizeof(c_uint) * c_count.value # NOTE: use 'c_uint' to represent the size of the nvmlPstate_t enumeration. pstates_array = _nvmlPstates_t * c_count.value @@ -5430,6 +5825,7 @@ def nvmlDeviceGetSupportedPerformanceStates(device): return pstates + def nvmlDeviceGetGpcClkVfOffset(device): offset = c_int32() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkVfOffset") @@ -5437,6 +5833,7 @@ def nvmlDeviceGetGpcClkVfOffset(device): _nvmlCheckReturn(ret) return offset.value + def nvmlDeviceSetGpcClkVfOffset(device, offset): c_offset = c_int32(offset) fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpcClkVfOffset") @@ -5444,6 +5841,7 @@ def nvmlDeviceSetGpcClkVfOffset(device, offset): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetGpcClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()): isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int) minOffsetRef = minOffset if isReference else byref(minOffset) @@ -5453,6 +5851,7 @@ def nvmlDeviceGetGpcClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) + def nvmlDeviceGetMemClkVfOffset(device): offset = c_int32() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkVfOffset") @@ -5460,6 +5859,7 @@ def nvmlDeviceGetMemClkVfOffset(device): _nvmlCheckReturn(ret) return offset.value + def nvmlDeviceSetMemClkVfOffset(device, offset): c_offset = c_int32(offset) fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemClkVfOffset") @@ -5467,6 +5867,7 @@ def nvmlDeviceSetMemClkVfOffset(device, offset): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetMemClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()): isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int) minOffsetRef = minOffset if isReference else byref(minOffset) @@ -5477,6 +5878,7 @@ def nvmlDeviceGetMemClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) + def nvmlSystemSetConfComputeGpusReadyState(state): c_state = c_uint(state) fn = _nvmlGetFunctionPointer("nvmlSystemSetConfComputeGpusReadyState") @@ -5484,6 +5886,7 @@ def nvmlSystemSetConfComputeGpusReadyState(state): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlSystemGetConfComputeGpusReadyState(): c_state = c_uint() fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeGpusReadyState") @@ -5491,6 +5894,7 @@ def nvmlSystemGetConfComputeGpusReadyState(): _nvmlCheckReturn(ret) return c_state.value + def nvmlSystemGetConfComputeCapabilities(): c_ccSysCaps = c_nvmlConfComputeSystemCaps_t() fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeCapabilities") @@ -5498,6 +5902,7 @@ def nvmlSystemGetConfComputeCapabilities(): _nvmlCheckReturn(ret) return c_ccSysCaps + def nvmlSystemGetConfComputeState(): c_state = c_nvmlConfComputeSystemState_t() fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeState") @@ -5505,16 +5910,19 @@ def nvmlSystemGetConfComputeState(): _nvmlCheckReturn(ret) return c_state + def nvmlSystemGetConfComputeSettings(settings): fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeSettings") return fn(settings) + def nvmlDeviceSetConfComputeUnprotectedMemSize(device, c_ccMemSize): fn = _nvmlGetFunctionPointer("nvmlDeviceSetConfComputeUnprotectedMemSize") ret = fn(device, c_ccMemSize) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetConfComputeMemSizeInfo(device): c_ccMemSize = c_nvmlConfComputeMemSizeInfo_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeMemSizeInfo") @@ -5522,6 +5930,7 @@ def nvmlDeviceGetConfComputeMemSizeInfo(device): _nvmlCheckReturn(ret) return c_ccMemSize + def nvmlDeviceGetConfComputeProtectedMemoryUsage(device): c_memory = c_nvmlMemory_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeProtectedMemoryUsage") @@ -5529,6 +5938,7 @@ def nvmlDeviceGetConfComputeProtectedMemoryUsage(device): _nvmlCheckReturn(ret) return c_memory + def nvmlDeviceGetConfComputeGpuCertificate(device): c_cert = c_nvmlConfComputeGpuCertificate_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuCertificate") @@ -5536,6 +5946,7 @@ def nvmlDeviceGetConfComputeGpuCertificate(device): _nvmlCheckReturn(ret) return c_cert + def nvmlDeviceGetConfComputeGpuAttestationReport(device, c_nonce): c_attestReport = c_nvmlConfComputeGpuAttestationReport_t() c_nonce_arr = (c_uint8 * len(c_nonce))(*(c_nonce)) @@ -5545,6 +5956,7 @@ def nvmlDeviceGetConfComputeGpuAttestationReport(device, c_nonce): _nvmlCheckReturn(ret) return c_attestReport + def nvmlSystemSetConfComputeKeyRotationThresholdInfo(max_atk_adv): c_keyRotationThrInfo = c_nvmlConfComputeSetKeyRotationThresholdInfo_t(0) c_keyRotationThrInfo.version = ConfComputeSetKeyRotationThresholdInfo_v1 @@ -5554,6 +5966,7 @@ def nvmlSystemSetConfComputeKeyRotationThresholdInfo(max_atk_adv): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlSystemGetConfComputeKeyRotationThresholdInfo(): c_keyRotationThrInfo = c_nvmlConfComputeGetKeyRotationThresholdInfo_t(0) c_keyRotationThrInfo.version = ConfComputeGetKeyRotationThresholdInfo_v1 @@ -5562,86 +5975,88 @@ def nvmlSystemGetConfComputeKeyRotationThresholdInfo(): _nvmlCheckReturn(ret) return c_keyRotationThrInfo + ## GPM ## ######### ## Enums/defines #### GPM Metric Identifiers -NVML_GPM_METRIC_GRAPHICS_UTIL = 1 # Percentage of time any compute/graphics app was active on the GPU. 0.0 - 100.0 -NVML_GPM_METRIC_SM_UTIL = 2 # Percentage of SMs that were busy. 0.0 - 100.0 -NVML_GPM_METRIC_SM_OCCUPANCY = 3 # Percentage of warps that were active vs theoretical maximum. 0.0 - 100.0 -NVML_GPM_METRIC_INTEGER_UTIL = 4 # Percentage of time the GPU's SMs were doing integer operations. 0.0 - 100.0 -NVML_GPM_METRIC_ANY_TENSOR_UTIL = 5 # Percentage of time the GPU's SMs were doing ANY tensor operations. 0.0 - 100.0 -NVML_GPM_METRIC_DFMA_TENSOR_UTIL = 6 # Percentage of time the GPU's SMs were doing DFMA tensor operations. 0.0 - 100.0 -NVML_GPM_METRIC_HMMA_TENSOR_UTIL = 7 # Percentage of time the GPU's SMs were doing HMMA tensor operations. 0.0 - 100.0 -NVML_GPM_METRIC_IMMA_TENSOR_UTIL = 9 # Percentage of time the GPU's SMs were doing IMMA tensor operations. 0.0 - 100.0 -NVML_GPM_METRIC_DRAM_BW_UTIL = 10 # Percentage of DRAM bw used vs theoretical maximum. 0.0 - 100.0 -NVML_GPM_METRIC_FP64_UTIL = 11 # Percentage of time the GPU's SMs were doing non-tensor FP64 math. 0.0 - 100.0 -NVML_GPM_METRIC_FP32_UTIL = 12 # Percentage of time the GPU's SMs were doing non-tensor FP32 math. 0.0 - 100.0 -NVML_GPM_METRIC_FP16_UTIL = 13 # Percentage of time the GPU's SMs were doing non-tensor FP16 math. 0.0 - 100.0 -NVML_GPM_METRIC_PCIE_TX_PER_SEC = 20 # PCIe traffic from this GPU in MiB/sec -NVML_GPM_METRIC_PCIE_RX_PER_SEC = 21 # PCIe traffic to this GPU in MiB/sec -NVML_GPM_METRIC_NVDEC_0_UTIL = 30 # Percent utilization of NVDEC 0. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_1_UTIL = 31 # Percent utilization of NVDEC 1. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_2_UTIL = 32 # Percent utilization of NVDEC 2. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_3_UTIL = 33 # Percent utilization of NVDEC 3. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_4_UTIL = 34 # Percent utilization of NVDEC 4. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_5_UTIL = 35 # Percent utilization of NVDEC 5. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_6_UTIL = 36 # Percent utilization of NVDEC 6. 0.0 - 100.0 -NVML_GPM_METRIC_NVDEC_7_UTIL = 37 # Percent utilization of NVDEC 7. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_0_UTIL = 40 # Percent utilization of NVJPG 0. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_1_UTIL = 41 # Percent utilization of NVJPG 1. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_2_UTIL = 42 # Percent utilization of NVJPG 2. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_3_UTIL = 43 # Percent utilization of NVJPG 3. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_4_UTIL = 44 # Percent utilization of NVJPG 4. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_5_UTIL = 45 # Percent utilization of NVJPG 5. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_6_UTIL = 46 # Percent utilization of NVJPG 6. 0.0 - 100.0 -NVML_GPM_METRIC_NVJPG_7_UTIL = 47 # Percent utilization of NVJPG 7. 0.0 - 100.0 -NVML_GPM_METRIC_NVOFA_0_UTIL = 50 # Percent utilization of NVOFA 0. 0.0 - 100.0 -NVML_GPM_METRIC_NVOFA_1_UTIL = 51 # Percent utilization of NVOFA 1. 0.0 - 100.0 -NVML_GPM_METRIC_NVLINK_TOTAL_RX_PER_SEC = 60 # NvLink read bandwidth for all links in MiB/sec -NVML_GPM_METRIC_NVLINK_TOTAL_TX_PER_SEC = 61 # NvLink write bandwidth for all links in MiB/sec -NVML_GPM_METRIC_NVLINK_L0_RX_PER_SEC = 62 # NvLink read bandwidth for link 0 in MiB/sec -NVML_GPM_METRIC_NVLINK_L0_TX_PER_SEC = 63 # NvLink write bandwidth for link 0 in MiB/sec -NVML_GPM_METRIC_NVLINK_L1_RX_PER_SEC = 64 # NvLink read bandwidth for link 1 in MiB/sec -NVML_GPM_METRIC_NVLINK_L1_TX_PER_SEC = 65 # NvLink write bandwidth for link 1 in MiB/sec -NVML_GPM_METRIC_NVLINK_L2_RX_PER_SEC = 66 # NvLink read bandwidth for link 2 in MiB/sec -NVML_GPM_METRIC_NVLINK_L2_TX_PER_SEC = 67 # NvLink write bandwidth for link 2 in MiB/sec -NVML_GPM_METRIC_NVLINK_L3_RX_PER_SEC = 68 # NvLink read bandwidth for link 3 in MiB/sec -NVML_GPM_METRIC_NVLINK_L3_TX_PER_SEC = 69 # NvLink write bandwidth for link 3 in MiB/sec -NVML_GPM_METRIC_NVLINK_L4_RX_PER_SEC = 70 # NvLink read bandwidth for link 4 in MiB/sec -NVML_GPM_METRIC_NVLINK_L4_TX_PER_SEC = 71 # NvLink write bandwidth for link 4 in MiB/sec -NVML_GPM_METRIC_NVLINK_L5_RX_PER_SEC = 72 # NvLink read bandwidth for link 5 in MiB/sec -NVML_GPM_METRIC_NVLINK_L5_TX_PER_SEC = 73 # NvLink write bandwidth for link 5 in MiB/sec -NVML_GPM_METRIC_NVLINK_L6_RX_PER_SEC = 74 # NvLink read bandwidth for link 6 in MiB/sec -NVML_GPM_METRIC_NVLINK_L6_TX_PER_SEC = 75 # NvLink write bandwidth for link 6 in MiB/sec -NVML_GPM_METRIC_NVLINK_L7_RX_PER_SEC = 76 # NvLink read bandwidth for link 7 in MiB/sec -NVML_GPM_METRIC_NVLINK_L7_TX_PER_SEC = 77 # NvLink write bandwidth for link 7 in MiB/sec -NVML_GPM_METRIC_NVLINK_L8_RX_PER_SEC = 78 # NvLink read bandwidth for link 8 in MiB/sec -NVML_GPM_METRIC_NVLINK_L8_TX_PER_SEC = 79 # NvLink write bandwidth for link 8 in MiB/sec -NVML_GPM_METRIC_NVLINK_L9_RX_PER_SEC = 80 # NvLink read bandwidth for link 9 in MiB/sec -NVML_GPM_METRIC_NVLINK_L9_TX_PER_SEC = 81 # NvLink write bandwidth for link 9 in MiB/sec -NVML_GPM_METRIC_NVLINK_L10_RX_PER_SEC = 82 # NvLink read bandwidth for link 10 in MiB/sec -NVML_GPM_METRIC_NVLINK_L10_TX_PER_SEC = 83 # NvLink write bandwidth for link 10 in MiB/sec -NVML_GPM_METRIC_NVLINK_L11_RX_PER_SEC = 84 # NvLink read bandwidth for link 11 in MiB/sec -NVML_GPM_METRIC_NVLINK_L11_TX_PER_SEC = 85 # NvLink write bandwidth for link 11 in MiB/sec -NVML_GPM_METRIC_NVLINK_L12_RX_PER_SEC = 86 # NvLink read bandwidth for link 12 in MiB/sec -NVML_GPM_METRIC_NVLINK_L12_TX_PER_SEC = 87 # NvLink write bandwidth for link 12 in MiB/sec -NVML_GPM_METRIC_NVLINK_L13_RX_PER_SEC = 88 # NvLink read bandwidth for link 13 in MiB/sec -NVML_GPM_METRIC_NVLINK_L13_TX_PER_SEC = 89 # NvLink write bandwidth for link 13 in MiB/sec -NVML_GPM_METRIC_NVLINK_L14_RX_PER_SEC = 90 # NvLink read bandwidth for link 14 in MiB/sec -NVML_GPM_METRIC_NVLINK_L14_TX_PER_SEC = 91 # NvLink write bandwidth for link 14 in MiB/sec -NVML_GPM_METRIC_NVLINK_L15_RX_PER_SEC = 92 # NvLink read bandwidth for link 15 in MiB/sec -NVML_GPM_METRIC_NVLINK_L15_TX_PER_SEC = 93 # NvLink write bandwidth for link 15 in MiB/sec -NVML_GPM_METRIC_NVLINK_L16_RX_PER_SEC = 94 # NvLink read bandwidth for link 16 in MiB/sec -NVML_GPM_METRIC_NVLINK_L16_TX_PER_SEC = 95 # NvLink write bandwidth for link 16 in MiB/sec -NVML_GPM_METRIC_NVLINK_L17_RX_PER_SEC = 96 # NvLink read bandwidth for link 17 in MiB/sec -NVML_GPM_METRIC_NVLINK_L17_TX_PER_SEC = 97 # NvLink write bandwidth for link 17 in MiB/sec -NVML_GPM_METRIC_MAX = 98 +NVML_GPM_METRIC_GRAPHICS_UTIL = 1 # Percentage of time any compute/graphics app was active on the GPU. 0.0 - 100.0 +NVML_GPM_METRIC_SM_UTIL = 2 # Percentage of SMs that were busy. 0.0 - 100.0 +NVML_GPM_METRIC_SM_OCCUPANCY = 3 # Percentage of warps that were active vs theoretical maximum. 0.0 - 100.0 +NVML_GPM_METRIC_INTEGER_UTIL = 4 # Percentage of time the GPU's SMs were doing integer operations. 0.0 - 100.0 +NVML_GPM_METRIC_ANY_TENSOR_UTIL = 5 # Percentage of time the GPU's SMs were doing ANY tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_DFMA_TENSOR_UTIL = 6 # Percentage of time the GPU's SMs were doing DFMA tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_HMMA_TENSOR_UTIL = 7 # Percentage of time the GPU's SMs were doing HMMA tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_IMMA_TENSOR_UTIL = 9 # Percentage of time the GPU's SMs were doing IMMA tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_DRAM_BW_UTIL = 10 # Percentage of DRAM bw used vs theoretical maximum. 0.0 - 100.0 +NVML_GPM_METRIC_FP64_UTIL = 11 # Percentage of time the GPU's SMs were doing non-tensor FP64 math. 0.0 - 100.0 +NVML_GPM_METRIC_FP32_UTIL = 12 # Percentage of time the GPU's SMs were doing non-tensor FP32 math. 0.0 - 100.0 +NVML_GPM_METRIC_FP16_UTIL = 13 # Percentage of time the GPU's SMs were doing non-tensor FP16 math. 0.0 - 100.0 +NVML_GPM_METRIC_PCIE_TX_PER_SEC = 20 # PCIe traffic from this GPU in MiB/sec +NVML_GPM_METRIC_PCIE_RX_PER_SEC = 21 # PCIe traffic to this GPU in MiB/sec +NVML_GPM_METRIC_NVDEC_0_UTIL = 30 # Percent utilization of NVDEC 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_1_UTIL = 31 # Percent utilization of NVDEC 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_2_UTIL = 32 # Percent utilization of NVDEC 2. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_3_UTIL = 33 # Percent utilization of NVDEC 3. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_4_UTIL = 34 # Percent utilization of NVDEC 4. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_5_UTIL = 35 # Percent utilization of NVDEC 5. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_6_UTIL = 36 # Percent utilization of NVDEC 6. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_7_UTIL = 37 # Percent utilization of NVDEC 7. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_0_UTIL = 40 # Percent utilization of NVJPG 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_1_UTIL = 41 # Percent utilization of NVJPG 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_2_UTIL = 42 # Percent utilization of NVJPG 2. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_3_UTIL = 43 # Percent utilization of NVJPG 3. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_4_UTIL = 44 # Percent utilization of NVJPG 4. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_5_UTIL = 45 # Percent utilization of NVJPG 5. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_6_UTIL = 46 # Percent utilization of NVJPG 6. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_7_UTIL = 47 # Percent utilization of NVJPG 7. 0.0 - 100.0 +NVML_GPM_METRIC_NVOFA_0_UTIL = 50 # Percent utilization of NVOFA 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVOFA_1_UTIL = 51 # Percent utilization of NVOFA 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVLINK_TOTAL_RX_PER_SEC = 60 # NvLink read bandwidth for all links in MiB/sec +NVML_GPM_METRIC_NVLINK_TOTAL_TX_PER_SEC = 61 # NvLink write bandwidth for all links in MiB/sec +NVML_GPM_METRIC_NVLINK_L0_RX_PER_SEC = 62 # NvLink read bandwidth for link 0 in MiB/sec +NVML_GPM_METRIC_NVLINK_L0_TX_PER_SEC = 63 # NvLink write bandwidth for link 0 in MiB/sec +NVML_GPM_METRIC_NVLINK_L1_RX_PER_SEC = 64 # NvLink read bandwidth for link 1 in MiB/sec +NVML_GPM_METRIC_NVLINK_L1_TX_PER_SEC = 65 # NvLink write bandwidth for link 1 in MiB/sec +NVML_GPM_METRIC_NVLINK_L2_RX_PER_SEC = 66 # NvLink read bandwidth for link 2 in MiB/sec +NVML_GPM_METRIC_NVLINK_L2_TX_PER_SEC = 67 # NvLink write bandwidth for link 2 in MiB/sec +NVML_GPM_METRIC_NVLINK_L3_RX_PER_SEC = 68 # NvLink read bandwidth for link 3 in MiB/sec +NVML_GPM_METRIC_NVLINK_L3_TX_PER_SEC = 69 # NvLink write bandwidth for link 3 in MiB/sec +NVML_GPM_METRIC_NVLINK_L4_RX_PER_SEC = 70 # NvLink read bandwidth for link 4 in MiB/sec +NVML_GPM_METRIC_NVLINK_L4_TX_PER_SEC = 71 # NvLink write bandwidth for link 4 in MiB/sec +NVML_GPM_METRIC_NVLINK_L5_RX_PER_SEC = 72 # NvLink read bandwidth for link 5 in MiB/sec +NVML_GPM_METRIC_NVLINK_L5_TX_PER_SEC = 73 # NvLink write bandwidth for link 5 in MiB/sec +NVML_GPM_METRIC_NVLINK_L6_RX_PER_SEC = 74 # NvLink read bandwidth for link 6 in MiB/sec +NVML_GPM_METRIC_NVLINK_L6_TX_PER_SEC = 75 # NvLink write bandwidth for link 6 in MiB/sec +NVML_GPM_METRIC_NVLINK_L7_RX_PER_SEC = 76 # NvLink read bandwidth for link 7 in MiB/sec +NVML_GPM_METRIC_NVLINK_L7_TX_PER_SEC = 77 # NvLink write bandwidth for link 7 in MiB/sec +NVML_GPM_METRIC_NVLINK_L8_RX_PER_SEC = 78 # NvLink read bandwidth for link 8 in MiB/sec +NVML_GPM_METRIC_NVLINK_L8_TX_PER_SEC = 79 # NvLink write bandwidth for link 8 in MiB/sec +NVML_GPM_METRIC_NVLINK_L9_RX_PER_SEC = 80 # NvLink read bandwidth for link 9 in MiB/sec +NVML_GPM_METRIC_NVLINK_L9_TX_PER_SEC = 81 # NvLink write bandwidth for link 9 in MiB/sec +NVML_GPM_METRIC_NVLINK_L10_RX_PER_SEC = 82 # NvLink read bandwidth for link 10 in MiB/sec +NVML_GPM_METRIC_NVLINK_L10_TX_PER_SEC = 83 # NvLink write bandwidth for link 10 in MiB/sec +NVML_GPM_METRIC_NVLINK_L11_RX_PER_SEC = 84 # NvLink read bandwidth for link 11 in MiB/sec +NVML_GPM_METRIC_NVLINK_L11_TX_PER_SEC = 85 # NvLink write bandwidth for link 11 in MiB/sec +NVML_GPM_METRIC_NVLINK_L12_RX_PER_SEC = 86 # NvLink read bandwidth for link 12 in MiB/sec +NVML_GPM_METRIC_NVLINK_L12_TX_PER_SEC = 87 # NvLink write bandwidth for link 12 in MiB/sec +NVML_GPM_METRIC_NVLINK_L13_RX_PER_SEC = 88 # NvLink read bandwidth for link 13 in MiB/sec +NVML_GPM_METRIC_NVLINK_L13_TX_PER_SEC = 89 # NvLink write bandwidth for link 13 in MiB/sec +NVML_GPM_METRIC_NVLINK_L14_RX_PER_SEC = 90 # NvLink read bandwidth for link 14 in MiB/sec +NVML_GPM_METRIC_NVLINK_L14_TX_PER_SEC = 91 # NvLink write bandwidth for link 14 in MiB/sec +NVML_GPM_METRIC_NVLINK_L15_RX_PER_SEC = 92 # NvLink read bandwidth for link 15 in MiB/sec +NVML_GPM_METRIC_NVLINK_L15_TX_PER_SEC = 93 # NvLink write bandwidth for link 15 in MiB/sec +NVML_GPM_METRIC_NVLINK_L16_RX_PER_SEC = 94 # NvLink read bandwidth for link 16 in MiB/sec +NVML_GPM_METRIC_NVLINK_L16_TX_PER_SEC = 95 # NvLink write bandwidth for link 16 in MiB/sec +NVML_GPM_METRIC_NVLINK_L17_RX_PER_SEC = 96 # NvLink read bandwidth for link 17 in MiB/sec +NVML_GPM_METRIC_NVLINK_L17_TX_PER_SEC = 97 # NvLink write bandwidth for link 17 in MiB/sec +NVML_GPM_METRIC_MAX = 98 ## Structs + class c_nvmlUnitInfo_t(_PrintableStructure): _fields_ = [ ('name', c_char * 96), @@ -5650,10 +6065,14 @@ class c_nvmlUnitInfo_t(_PrintableStructure): ('firmwareVersion', c_char * 96), ] + class struct_c_nvmlGpmSample_t(Structure): - pass # opaque handle + pass # opaque handle + + c_nvmlGpmSample_t = POINTER(struct_c_nvmlGpmSample_t) + class c_metricInfo_t(Structure): _fields_ = [ ("shortName", c_char_p), @@ -5661,47 +6080,46 @@ class c_metricInfo_t(Structure): ("unit", c_char_p), ] + class c_nvmlGpmMetric_t(_PrintableStructure): - _fields_ = [ - ('metricId', c_uint), - ('nvmlReturn', _nvmlReturn_t), - ('value', c_double), - ('metricInfo', c_metricInfo_t) - ] + _fields_ = [('metricId', c_uint), ('nvmlReturn', _nvmlReturn_t), ('value', c_double), + ('metricInfo', c_metricInfo_t)] + class c_nvmlGpmMetricsGet_t(_PrintableStructure): - _fields_ = [ - ('version', c_uint), - ('numMetrics', c_uint), - ('sample1', c_nvmlGpmSample_t), - ('sample2', c_nvmlGpmSample_t), - ('metrics', c_nvmlGpmMetric_t * NVML_GPM_METRIC_MAX) - ] + _fields_ = [('version', c_uint), ('numMetrics', c_uint), ('sample1', c_nvmlGpmSample_t), + ('sample2', c_nvmlGpmSample_t), ('metrics', c_nvmlGpmMetric_t * NVML_GPM_METRIC_MAX)] + NVML_GPM_METRICS_GET_VERSION = 1 + class c_nvmlGpmSupport_t(_PrintableStructure): _fields_ = [ ('version', c_uint), ('isSupportedDevice', c_uint), ] + NVML_GPM_SUPPORT_VERSION = 1 ## Functions + def nvmlGpmMetricsGet(metricsGet): fn = _nvmlGetFunctionPointer("nvmlGpmMetricsGet") ret = fn(byref(metricsGet)) _nvmlCheckReturn(ret) return metricsGet + def nvmlGpmSampleFree(gpmSample): fn = _nvmlGetFunctionPointer("nvmlGpmSampleFree") ret = fn(gpmSample) _nvmlCheckReturn(ret) return + def nvmlGpmSampleAlloc(): gpmSample = c_nvmlGpmSample_t() fn = _nvmlGetFunctionPointer("nvmlGpmSampleAlloc") @@ -5709,18 +6127,21 @@ def nvmlGpmSampleAlloc(): _nvmlCheckReturn(ret) return gpmSample + def nvmlGpmSampleGet(device, gpmSample): fn = _nvmlGetFunctionPointer("nvmlGpmSampleGet") ret = fn(device, gpmSample) _nvmlCheckReturn(ret) return gpmSample + def nvmlGpmMigSampleGet(device, gpuInstanceId, gpmSample): fn = _nvmlGetFunctionPointer("nvmlGpmMigSampleGet") ret = fn(device, gpuInstanceId, gpmSample) _nvmlCheckReturn(ret) return gpmSample + def nvmlGpmQueryDeviceSupport(device): gpmSupport = c_nvmlGpmSupport_t() gpmSupport.version = NVML_GPM_SUPPORT_VERSION @@ -5729,6 +6150,7 @@ def nvmlGpmQueryDeviceSupport(device): _nvmlCheckReturn(ret) return gpmSupport + def nvmlGpmSetStreamingEnabled(device, state): c_state = c_uint(state) fn = _nvmlGetFunctionPointer("nvmlGpmSetStreamingEnabled") @@ -5736,6 +6158,7 @@ def nvmlGpmSetStreamingEnabled(device, state): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlGpmQueryIfStreamingEnabled(device): c_state = c_uint() fn = _nvmlGetFunctionPointer("nvmlGpmQueryIfStreamingEnabled") @@ -5743,21 +6166,24 @@ def nvmlGpmQueryIfStreamingEnabled(device): _nvmlCheckReturn(ret) return c_state.value + # Low Power Structure and Function -NVML_NVLINK_POWER_STATE_HIGH_SPEED = 0x0 -NVML_NVLINK_POWER_STATE_LOW = 0x1 +NVML_NVLINK_POWER_STATE_HIGH_SPEED = 0x0 +NVML_NVLINK_POWER_STATE_LOW = 0x1 -NVML_NVLINK_LOW_POWER_THRESHOLD_MIN = 0x1 -NVML_NVLINK_LOW_POWER_THRESHOLD_MAX = 0x1FFF -NVML_NVLINK_LOW_POWER_THRESHOLD_RESET = 0xFFFFFFFF +NVML_NVLINK_LOW_POWER_THRESHOLD_MIN = 0x1 +NVML_NVLINK_LOW_POWER_THRESHOLD_MAX = 0x1FFF +NVML_NVLINK_LOW_POWER_THRESHOLD_RESET = 0xFFFFFFFF NVML_NVLINK_LOW_POWER_THRESHOLD_DEFAULT = NVML_NVLINK_LOW_POWER_THRESHOLD_RESET + class c_nvmlNvLinkPowerThres_t(Structure): _fields_ = [ ("lowPwrThreshold", c_uint), ] + def nvmlDeviceSetNvLinkDeviceLowPowerThreshold(device, l1threshold): c_info = c_nvmlNvLinkPowerThres_t() c_info.lowPwrThreshold = l1threshold @@ -5766,82 +6192,80 @@ def nvmlDeviceSetNvLinkDeviceLowPowerThreshold(device, l1threshold): _nvmlCheckReturn(ret) return NVML_SUCCESS + NVML_GPU_FABRIC_UUID_LEN = 16 _nvmlGpuFabricState_t = c_uint NVML_GPU_FABRIC_STATE_NOT_SUPPORTED = 0 -NVML_GPU_FABRIC_STATE_NOT_STARTED = 1 -NVML_GPU_FABRIC_STATE_IN_PROGRESS = 2 -NVML_GPU_FABRIC_STATE_COMPLETED = 3 +NVML_GPU_FABRIC_STATE_NOT_STARTED = 1 +NVML_GPU_FABRIC_STATE_IN_PROGRESS = 2 +NVML_GPU_FABRIC_STATE_COMPLETED = 3 + class c_nvmlGpuFabricInfo_t(_PrintableStructure): - _fields_ = [ - ("clusterUuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE), - ("status", _nvmlReturn_t), - ("cliqueId", c_uint32), - ("state", _nvmlGpuFabricState_t) - ] + _fields_ = [("clusterUuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE), ("status", _nvmlReturn_t), + ("cliqueId", c_uint32), ("state", _nvmlGpuFabricState_t)] + NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_NOT_SUPPORTED = 0 -NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_TRUE = 1 -NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_FALSE = 2 -NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_DEGRADED_BW = 0 -NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_DEGRADED_BW = 0x11 +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_DEGRADED_BW = 0 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_DEGRADED_BW = 0x11 -NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_NOT_SUPPORTED = 0 -NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_TRUE = 1 -NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_FALSE = 2 -NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_RECOVERY = 2 -NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_RECOVERY = 0x11 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_RECOVERY = 2 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_RECOVERY = 0x11 -NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_NOT_SUPPORTED = 0 -NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_TRUE = 1 -NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_FALSE = 2 -NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_UNHEALTHY = 4 -NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_UNHEALTHY = 0x11 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_UNHEALTHY = 4 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_UNHEALTHY = 0x11 NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_NOT_SUPPORTED = 0 -NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_TRUE = 1 -NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_FALSE = 2 -NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ACCESS_TIMEOUT_RECOVERY = 6 -NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ACCESS_TIMEOUT_RECOVERY = 0x11 +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ACCESS_TIMEOUT_RECOVERY = 6 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ACCESS_TIMEOUT_RECOVERY = 0x11 nvmlGpuFabricInfo_v2 = 0x02000024 + class c_nvmlGpuFabricInfoV_t(_PrintableStructure): - _fields_ = [ - ("version", c_uint), - ("clusterUuid", c_char * NVML_GPU_FABRIC_UUID_LEN), - ("status", _nvmlReturn_t), - ("cliqueId", c_uint32), - ("state", _nvmlGpuFabricState_t), - ("healthMask", c_uint32) - ] + _fields_ = [("version", c_uint), ("clusterUuid", c_char * NVML_GPU_FABRIC_UUID_LEN), ("status", _nvmlReturn_t), + ("cliqueId", c_uint32), ("state", _nvmlGpuFabricState_t), ("healthMask", c_uint32)] def __init__(self): super(c_nvmlGpuFabricInfoV_t, self).__init__(version=nvmlGpuFabricInfo_v2) + def nvmlDeviceGetGpuFabricInfo(device, gpuFabricInfo): - fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfo"); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfo") ret = fn(device, gpuFabricInfo) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetGpuFabricInfoV(device, gpuFabricInfo): - fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfoV"); + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfoV") ret = fn(device, gpuFabricInfo) _nvmlCheckReturn(ret) return NVML_SUCCESS + ###################### ## Enums/defines #### NVML GPU NVLINK BW MODE -NVML_GPU_NVLINK_BW_MODE_FULL = 0x0 -NVML_GPU_NVLINK_BW_MODE_OFF = 0x1 -NVML_GPU_NVLINK_BW_MODE_MIN = 0x2 -NVML_GPU_NVLINK_BW_MODE_HALF = 0x3 -NVML_GPU_NVLINK_BW_MODE_3QUARTER = 0x4 -NVML_GPU_NVLINK_BW_MODE_COUNT = 0x5 +NVML_GPU_NVLINK_BW_MODE_FULL = 0x0 +NVML_GPU_NVLINK_BW_MODE_OFF = 0x1 +NVML_GPU_NVLINK_BW_MODE_MIN = 0x2 +NVML_GPU_NVLINK_BW_MODE_HALF = 0x3 +NVML_GPU_NVLINK_BW_MODE_3QUARTER = 0x4 +NVML_GPU_NVLINK_BW_MODE_COUNT = 0x5 + def nvmlSystemSetNvlinkBwMode(mode): fn = _nvmlGetFunctionPointer("nvmlSystemSetNvlinkBwMode") @@ -5849,6 +6273,7 @@ def nvmlSystemSetNvlinkBwMode(mode): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlSystemGetNvlinkBwMode(): mode = c_uint() fn = _nvmlGetFunctionPointer("nvmlSystemGetNvlinkBwMode") @@ -5856,10 +6281,12 @@ def nvmlSystemGetNvlinkBwMode(): _nvmlCheckReturn(ret) return mode.value + _nvmlPowerScopeType_t = c_uint -NVML_POWER_SCOPE_GPU = 0 -NVML_POWER_SCOPE_MODULE = 1 -NVML_POWER_SCOPE_MEMORY = 2 +NVML_POWER_SCOPE_GPU = 0 +NVML_POWER_SCOPE_MODULE = 1 +NVML_POWER_SCOPE_MEMORY = 2 + class c_nvmlPowerValue_v2_t(_PrintableStructure): _fields_ = [ @@ -5869,8 +6296,10 @@ class c_nvmlPowerValue_v2_t(_PrintableStructure): ] _fmt_ = {'': "%d B"} + nvmlPowerValue_v2 = 0x0200000C + def nvmlDeviceSetPowerManagementLimit_v2(device, powerScope, powerLimit, version=nvmlPowerValue_v2): c_powerScope = _nvmlPowerScopeType_t(powerScope) c_powerValue = c_nvmlPowerValue_v2_t() @@ -5881,36 +6310,33 @@ def nvmlDeviceSetPowerManagementLimit_v2(device, powerScope, powerLimit, version ret = fn(device, byref(c_powerValue)) return NVML_SUCCESS + class c_nvmlEccSramErrorStatus_v1_t(_PrintableStructure): - _fields_ = [ - ('version', c_uint), - ('aggregateUncParity', c_ulonglong), - ('aggregateUncSecDed', c_ulonglong), - ('aggregateCor', c_ulonglong), - ('volatileUncParity', c_ulonglong), - ('volatileUncSecDed', c_ulonglong), - ('volatileCor', c_ulonglong), - ('aggregateUncBucketL2', c_ulonglong), - ('aggregateUncBucketSm', c_ulonglong), - ('aggregateUncBucketPcie', c_ulonglong), - ('aggregateUncBucketMcu', c_ulonglong), - ('aggregateUncBucketOther', c_ulonglong), - ('bThresholdExceeded', c_uint) - ] + _fields_ = [('version', c_uint), ('aggregateUncParity', c_ulonglong), ('aggregateUncSecDed', c_ulonglong), + ('aggregateCor', c_ulonglong), ('volatileUncParity', c_ulonglong), ('volatileUncSecDed', c_ulonglong), + ('volatileCor', c_ulonglong), ('aggregateUncBucketL2', c_ulonglong), + ('aggregateUncBucketSm', c_ulonglong), ('aggregateUncBucketPcie', c_ulonglong), + ('aggregateUncBucketMcu', c_ulonglong), ('aggregateUncBucketOther', c_ulonglong), + ('bThresholdExceeded', c_uint)] def __init__(self): super(c_nvmlEccSramErrorStatus_v1_t, self).__init__(version=nvmlEccSramErrorStatus_v1) + nvmlEccSramErrorStatus_v1 = 0x1000068 + + def nvmlDeviceGetSramEccErrorStatus(device, status): fn = _nvmlGetFunctionPointer("nvmlDeviceGetSramEccErrorStatus") ret = fn(device, status) _nvmlCheckReturn(ret) return NVML_SUCCESS + NVML_DEV_CAP_EGM = (1 << 0) nvmlDeviceCapabilities_v1 = 0x1000008 + class c_nvmlDeviceCapabilities_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -5925,87 +6351,87 @@ def nvmlDeviceGetCapabilities(device, caps): fn = _nvmlGetFunctionPointer("nvmlDeviceGetCapabilities") return fn(device, caps) + class c_nvmlPlatformInfo_v1_t(_PrintableStructure): - _fields_ = [ - ('version', c_uint), - ('ibGuid', c_char * 16), - ('rackGuid', c_char * 16), - ('chassisPhysicalSlotNumber', c_char), - ('computeSlotIndex', c_char), - ('nodeIndex', c_char), - ('peerType', c_char), - ('moduleId', c_char) - ] + _fields_ = [('version', c_uint), ('ibGuid', c_char * 16), ('rackGuid', c_char * 16), + ('chassisPhysicalSlotNumber', c_char), ('computeSlotIndex', c_char), ('nodeIndex', c_char), + ('peerType', c_char), ('moduleId', c_char)] def __init__(self): super(c_nvmlPlatformInfo_v1_t, self).__init__(version=nvmlPlatformInfo_v1) + nvmlPlatformInfo_v1 = 0x100002c + + def nvmlDeviceGetPlatformInfo(device, platformInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceGetPlatformInfo") ret = fn(device, platformInfo) _nvmlCheckReturn(ret) return NVML_SUCCESS + class c_nvmlMask255_t(_PrintableStructure): _fields_ = [ ('mask', c_uint * 8), ] -NVML_WORKLOAD_POWER_MAX_PROFILES = 255 -NVML_POWER_PROFILE_MAX_P = 0 -NVML_POWER_PROFILE_MAX_Q = 1 -NVML_POWER_PROFILE_COMPUTE = 2 -NVML_POWER_PROFILE_MEMORY_BOUND = 3 -NVML_POWER_PROFILE_NETWORK = 4 -NVML_POWER_PROFILE_BALANCED = 5 -NVML_POWER_PROFILE_LLM_INFERENCE = 6 -NVML_POWER_PROFILE_LLM_TRAINING = 7 -NVML_POWER_PROFILE_RBM = 8 -NVML_POWER_PROFILE_DCPCIE = 9 -NVML_POWER_PROFILE_HMMA_SPARSE = 10 -NVML_POWER_PROFILE_HMMA_DENSE = 11 -NVML_POWER_PROFILE_SYNC_BALANCED = 12 -NVML_POWER_PROFILE_HPC = 13 -NVML_POWER_PROFILE_MIG = 14 -NVML_POWER_PROFILE_MAX = 15 + +NVML_WORKLOAD_POWER_MAX_PROFILES = 255 +NVML_POWER_PROFILE_MAX_P = 0 +NVML_POWER_PROFILE_MAX_Q = 1 +NVML_POWER_PROFILE_COMPUTE = 2 +NVML_POWER_PROFILE_MEMORY_BOUND = 3 +NVML_POWER_PROFILE_NETWORK = 4 +NVML_POWER_PROFILE_BALANCED = 5 +NVML_POWER_PROFILE_LLM_INFERENCE = 6 +NVML_POWER_PROFILE_LLM_TRAINING = 7 +NVML_POWER_PROFILE_RBM = 8 +NVML_POWER_PROFILE_DCPCIE = 9 +NVML_POWER_PROFILE_HMMA_SPARSE = 10 +NVML_POWER_PROFILE_HMMA_DENSE = 11 +NVML_POWER_PROFILE_SYNC_BALANCED = 12 +NVML_POWER_PROFILE_HPC = 13 +NVML_POWER_PROFILE_MIG = 14 +NVML_POWER_PROFILE_MAX = 15 nvmlWorkloadPowerProfileInfo_v1 = 0x100002c + + class c_nvmlWorkloadPowerProfileInfo_v1_t(_PrintableStructure): - _fields_ = [ - ('version', c_uint), - ('profileId', c_uint), - ('priority', c_uint), - ('conflictingmask', c_nvmlMask255_t) - ] + _fields_ = [('version', c_uint), ('profileId', c_uint), ('priority', c_uint), ('conflictingmask', c_nvmlMask255_t)] def __init__(self): super(c_nvmlWorkloadPowerProfileInfo_v1_t, self).__init__(version=nvmlWorkloadPowerProfileInfo_v1) + nvmlWorkloadPowerProfileProfilesInfo_v1 = 0x1002bf8 + + class c_nvmlWorkloadPowerProfileProfilesInfo_v1_t(_PrintableStructure): - _fields_ = [ - ('version', c_uint), - ('perfProfilesMask', c_nvmlMask255_t), - ('perfProfile', c_nvmlWorkloadPowerProfileInfo_v1_t * NVML_WORKLOAD_POWER_MAX_PROFILES) - ] + _fields_ = [('version', c_uint), ('perfProfilesMask', c_nvmlMask255_t), + ('perfProfile', c_nvmlWorkloadPowerProfileInfo_v1_t * NVML_WORKLOAD_POWER_MAX_PROFILES)] def __init__(self): - super(c_nvmlWorkloadPowerProfileProfilesInfo_v1_t, self).__init__(version=nvmlWorkloadPowerProfileProfilesInfo_v1) + super(c_nvmlWorkloadPowerProfileProfilesInfo_v1_t, + self).__init__(version=nvmlWorkloadPowerProfileProfilesInfo_v1) + nvmlWorkloadPowerProfileCurrentProfiles_v1 = 0x1000064 + + class c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t(_PrintableStructure): - _fields_ = [ - ('version', c_uint), - ('perfProfilesMask', c_nvmlMask255_t), - ('requestedProfilesMask', c_nvmlMask255_t), - ('enforcedProfilesMask', c_nvmlMask255_t) - ] + _fields_ = [('version', c_uint), ('perfProfilesMask', c_nvmlMask255_t), ('requestedProfilesMask', c_nvmlMask255_t), + ('enforcedProfilesMask', c_nvmlMask255_t)] def __init__(self): - super(c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t, self).__init__(version=nvmlWorkloadPowerProfileCurrentProfiles_v1) + super(c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t, + self).__init__(version=nvmlWorkloadPowerProfileCurrentProfiles_v1) + nvmlWorkloadPowerProfileRequestedProfiles_v1 = 0x1000024 + + class c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -6013,7 +6439,9 @@ class c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t(_PrintableStructure): ] def __init__(self): - super(c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t, self).__init__(version=nvmlWorkloadPowerProfileRequestedProfiles_v1) + super(c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t, + self).__init__(version=nvmlWorkloadPowerProfileRequestedProfiles_v1) + def nvmlDeviceWorkloadPowerProfileGetProfilesInfo(device, profilesInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetProfilesInfo") @@ -6021,53 +6449,62 @@ def nvmlDeviceWorkloadPowerProfileGetProfilesInfo(device, profilesInfo): _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceWorkloadPowerProfileGetCurrentProfiles(device, currentProfiles): fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetCurrentProfiles") ret = fn(device, currentProfiles) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceWorkloadPowerProfileSetRequestedProfiles(device, requestedProfiles): fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileSetRequestedProfiles") ret = fn(device, requestedProfiles) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceWorkloadPowerProfileClearRequestedProfiles(device, requestedProfiles): fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileClearRequestedProfiles") ret = fn(device, requestedProfiles) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetNvlinkSupportedBwModes(device, supportedBwModes): fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkSupportedBwModes") ret = fn(device, supportedBwModes) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceGetNvlinkBwMode(device, getBwMode): fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkBwMode") ret = fn(device, getBwMode) _nvmlCheckReturn(ret) return NVML_SUCCESS + def nvmlDeviceSetNvlinkBwMode(device, setBwMode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvlinkBwMode") ret = fn(device, setBwMode) _nvmlCheckReturn(ret) return NVML_SUCCESS + nvmlDramEncryptionInfo_v1 = 0x01000008 + class c_nvmlDramEncryptionInfo_t(_PrintableStructure): _fields_ = [ ('version', c_uint), - ('encryptionState', _nvmlEnableState_t), + ('encryptionState', _nvmlEnableState_t), ] def __init__(self): super(c_nvmlDramEncryptionInfo_t, self).__init__(version=nvmlDramEncryptionInfo_v1) + def nvmlDeviceGetDramEncryptionMode(handle): c_currState = c_nvmlDramEncryptionInfo_t() c_pendingState = c_nvmlDramEncryptionInfo_t() @@ -6076,31 +6513,37 @@ def nvmlDeviceGetDramEncryptionMode(handle): _nvmlCheckReturn(ret) return [c_currState.encryptionState, c_pendingState.encryptionState] + # added to API def nvmlDeviceGetCurrentDramEncryptionMode(handle): return nvmlDeviceGetDramEncryptionMode(handle)[0] + # added to API def nvmlDeviceGetPendingDramEncryptionMode(handle): return nvmlDeviceGetDramEncryptionMode(handle)[1] + def nvmlDeviceSetDramEncryptionMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetDramEncryptionMode") c_dramEncryptionMode = c_nvmlDramEncryptionInfo_t() - c_dramEncryptionMode.encryptionState = mode; + c_dramEncryptionMode.encryptionState = mode ret = fn(handle, byref(c_dramEncryptionMode)) _nvmlCheckReturn(ret) return None + # Power Smoothing defines -NVML_POWER_SMOOTHING_MAX_NUM_PROFILES = 5 -NVML_POWER_SMOOTHING_ADMIN_OVERRIDE_NOT_SET = 0xFFFFFFFF -NVML_POWER_SMOOTHING_PROFILE_PARAM_PERCENT_TMP_FLOOR = 0 -NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_UP_RATE = 1 -NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_RATE = 2 +NVML_POWER_SMOOTHING_MAX_NUM_PROFILES = 5 +NVML_POWER_SMOOTHING_ADMIN_OVERRIDE_NOT_SET = 0xFFFFFFFF +NVML_POWER_SMOOTHING_PROFILE_PARAM_PERCENT_TMP_FLOOR = 0 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_UP_RATE = 1 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_RATE = 2 NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_HYSTERESIS = 3 -nvmlPowerSmoothingState_v1=0x1000008 +nvmlPowerSmoothingState_v1 = 0x1000008 + + class c_nvmlPowerSmoothingState_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -6110,7 +6553,10 @@ class c_nvmlPowerSmoothingState_v1_t(_PrintableStructure): def __init__(self): super(c_nvmlPowerSmoothingState_v1_t, self).__init__(version=nvmlPowerSmoothingState_v1) -nvmlPowerSmoothingProfile_v1=0x1000018 + +nvmlPowerSmoothingProfile_v1 = 0x1000018 + + class c_nvmlPowerSmoothingProfile_v1_t(_PrintableStructure): _fields_ = [ ('version', c_uint), @@ -6122,18 +6568,20 @@ class c_nvmlPowerSmoothingProfile_v1_t(_PrintableStructure): def __init__(self): super(c_nvmlPowerSmoothingProfile_v1_t, self).__init__(version=nvmlPowerSmoothingProfile_v1) + def nvmlDevicePowerSmoothingActivatePresetProfile(device, profile): fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingActivatePresetProfile") ret = fn(device, profile) _nvmlCheckReturn(ret) + def nvmlDevicePowerSmoothingUpdatePresetProfileParam(device, profile): fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingUpdatePresetProfileParam") ret = fn(device, profile) _nvmlCheckReturn(ret) + def nvmlDevicePowerSmoothingSetState(device, state): fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingSetState") ret = fn(device, state) _nvmlCheckReturn(ret) - diff --git a/scripts/check_docs_links.py b/scripts/check_docs_links.py index c80aa1618..898acec26 100644 --- a/scripts/check_docs_links.py +++ b/scripts/check_docs_links.py @@ -1,6 +1,5 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 - """Validate local markdown links under docs/.""" from __future__ import annotations @@ -12,8 +11,7 @@ LINK_PATTERN = re.compile(r"\[[^\]]+\]\(([^)]+)\)") def is_external(target: str) -> bool: - return target.startswith( - ("http://", "https://", "mailto:", "#", "<", "javascript:")) + return target.startswith(("http://", "https://", "mailto:", "#", "<", "javascript:")) def resolve_candidates(file_path: Path, docs_root: Path, target: str) -> list[Path]: @@ -56,9 +54,7 @@ def check_links(docs_root: Path) -> list[str]: if any(path.is_file() for path in candidates): continue - errors.append( - f"{md_file.relative_to(docs_root.parent)}:{line_no} -> {stripped}" - ) + errors.append(f"{md_file.relative_to(docs_root.parent)}:{line_no} -> {stripped}") return errors diff --git a/scripts/checkpoint_conversion/convert_flux2_klein.py b/scripts/checkpoint_conversion/convert_flux2_klein.py index 738dcc82d..2bf014ada 100644 --- a/scripts/checkpoint_conversion/convert_flux2_klein.py +++ b/scripts/checkpoint_conversion/convert_flux2_klein.py @@ -57,7 +57,6 @@ except ImportError: save_torch_state_dict = None snapshot_download = None - DEFAULT_REPO_ID = "black-forest-labs/FLUX.2-klein-4B" RAW_TRANSFORMER_FILENAME = "flux-2-klein-4b.safetensors" DIFFUSION_WEIGHTS_BASENAME = "diffusion_pytorch_model" @@ -165,8 +164,7 @@ def _resolve_src(src: str, revision: str | None, cache_dir: str | None, transfor cache_dir=cache_dir, token=_resolve_hf_token(), allow_patterns=_snapshot_allow_patterns(transformer_source), - ) - ) + )) def _read_json(path: Path) -> dict[str, Any]: @@ -312,9 +310,8 @@ def convert_raw_transformer(raw_state: OrderedDict[str, torch.Tensor]) -> Ordere unexpected.append(key) if unexpected: - raise ConversionError( - "Unmapped raw Flux2 Klein transformer keys:\n" + "\n".join(f" - {key}" for key in unexpected) - ) + raise ConversionError("Unmapped raw Flux2 Klein transformer keys:\n" + "\n".join(f" - {key}" + for key in unexpected)) return converted @@ -356,10 +353,7 @@ def _validate_transformer(state: OrderedDict[str, torch.Tensor]) -> None: keys = set(state) double_blocks = _infer_block_count(keys, "transformer_blocks.") single_blocks = _infer_block_count(keys, "single_transformer_blocks.") - print( - " transformer keys: " + - f"{len(state)} total, {double_blocks} double blocks, {single_blocks} single blocks" - ) + print(" transformer keys: " + f"{len(state)} total, {double_blocks} double blocks, {single_blocks} single blocks") def _validate_vae(state: OrderedDict[str, torch.Tensor]) -> None: @@ -403,18 +397,16 @@ def _build_model_index(src_dir: Path, source_label: str) -> dict[str, Any]: index = _read_json(index_path) else: index = dict(DEFAULT_MODEL_INDEX) - index.update( - { - "_class_name": "Flux2KleinPipeline", - "is_distilled": True, - "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], - "text_encoder": ["transformers", "Qwen3ForCausalLM"], - "tokenizer": ["transformers", "Qwen2TokenizerFast"], - "transformer": ["diffusers", "Flux2Transformer2DModel"], - "vae": ["diffusers", "AutoencoderKLFlux2"], - "_fastvideo_converted_from": source_label, - } - ) + index.update({ + "_class_name": "Flux2KleinPipeline", + "is_distilled": True, + "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], + "text_encoder": ["transformers", "Qwen3ForCausalLM"], + "tokenizer": ["transformers", "Qwen2TokenizerFast"], + "transformer": ["diffusers", "Flux2Transformer2DModel"], + "vae": ["diffusers", "AutoencoderKLFlux2"], + "_fastvideo_converted_from": source_label, + }) return index diff --git a/scripts/checkpoint_conversion/convert_gamecraft_full.py b/scripts/checkpoint_conversion/convert_gamecraft_full.py index 293480078..20057c26e 100644 --- a/scripts/checkpoint_conversion/convert_gamecraft_full.py +++ b/scripts/checkpoint_conversion/convert_gamecraft_full.py @@ -41,36 +41,15 @@ def create_model_index(output_dir: Path) -> None: model_index = { "_class_name": "HunyuanGameCraftPipeline", "_diffusers_version": "0.30.0", - "scheduler": [ - "diffusers", - "FlowMatchEulerDiscreteScheduler" - ], - "text_encoder": [ - "transformers", - "LlamaModel" - ], - "text_encoder_2": [ - "transformers", - "CLIPTextModel" - ], - "tokenizer": [ - "transformers", - "AutoTokenizer" - ], - "tokenizer_2": [ - "transformers", - "CLIPTokenizer" - ], - "transformer": [ - "fastvideo", - "HunyuanGameCraftTransformer3DModel" - ], - "vae": [ - "fastvideo", - "AutoencoderKLCausal3D" - ] + "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], + "text_encoder": ["transformers", "LlamaModel"], + "text_encoder_2": ["transformers", "CLIPTextModel"], + "tokenizer": ["transformers", "AutoTokenizer"], + "tokenizer_2": ["transformers", "CLIPTokenizer"], + "transformer": ["fastvideo", "HunyuanGameCraftTransformer3DModel"], + "vae": ["fastvideo", "AutoencoderKLCausal3D"] } - + with open(output_dir / "model_index.json", "w") as f: json.dump(model_index, f, indent=2) print(f"Created model_index.json") @@ -80,7 +59,7 @@ def create_scheduler_config(output_dir: Path) -> None: """Create scheduler config for FlowMatchEulerDiscreteScheduler.""" scheduler_dir = output_dir / "scheduler" scheduler_dir.mkdir(parents=True, exist_ok=True) - + scheduler_config = { "_class_name": "FlowMatchEulerDiscreteScheduler", "_diffusers_version": "0.30.0", @@ -92,7 +71,7 @@ def create_scheduler_config(output_dir: Path) -> None: "shift": 7.0, "use_dynamic_shifting": True } - + with open(scheduler_dir / "scheduler_config.json", "w") as f: json.dump(scheduler_config, f, indent=2) print(f"Created scheduler config") @@ -103,14 +82,14 @@ def copy_text_encoder(src_dir: Path, dst_dir: Path, name: str) -> None: if not src_dir.exists(): print(f"Warning: {name} not found at {src_dir}") return - + dst_dir.mkdir(parents=True, exist_ok=True) - + # Copy all files for src_file in src_dir.iterdir(): if src_file.is_file(): shutil.copy2(src_file, dst_dir / src_file.name) - + print(f"Copied {name} to {dst_dir}") @@ -119,9 +98,9 @@ def copy_tokenizer(src_dir: Path, dst_dir: Path, name: str) -> None: if not src_dir.exists(): print(f"Warning: {name} not found at {src_dir}") return - + dst_dir.mkdir(parents=True, exist_ok=True) - + # Copy tokenizer-related files tokenizer_files = [ "tokenizer.json", @@ -131,26 +110,26 @@ def copy_tokenizer(src_dir: Path, dst_dir: Path, name: str) -> None: "merges.txt", "added_tokens.json", ] - + for filename in tokenizer_files: src_file = src_dir / filename if src_file.exists(): shutil.copy2(src_file, dst_dir / filename) - + print(f"Copied {name} to {dst_dir}") def convert_full_model( - transformer_path: Path | None = None, - vae_path: Path | None = None, - text_encoder_path: Path | None = None, - text_encoder_2_path: Path | None = None, - input_dir: Path | None = None, - output_dir: Path = Path("official_weights/hunyuan-gamecraft-diffusers"), - verbose: bool = False, + transformer_path: Path | None = None, + vae_path: Path | None = None, + text_encoder_path: Path | None = None, + text_encoder_2_path: Path | None = None, + input_dir: Path | None = None, + output_dir: Path = Path("official_weights/hunyuan-gamecraft-diffusers"), + verbose: bool = False, ) -> None: """Convert all components of GameCraft to diffusers format.""" - + # Resolve paths from input_dir if individual paths not provided if input_dir is not None: if transformer_path is None: @@ -163,11 +142,11 @@ def convert_full_model( text_encoder_path = input_dir / "stdmodels" / "llava-llama-3-8b-v1_1-transformers" if text_encoder_2_path is None: text_encoder_2_path = input_dir / "stdmodels" / "openai_clip-vit-large-patch14" - + output_dir.mkdir(parents=True, exist_ok=True) print(f"Converting GameCraft to diffusers format at {output_dir}") print("=" * 60) - + # 1. Convert transformer if transformer_path and transformer_path.exists(): print(f"\n[1/5] Converting transformer from {transformer_path}") @@ -179,7 +158,7 @@ def convert_full_model( ) else: print(f"\n[1/5] Skipping transformer (not found: {transformer_path})") - + # 2. Convert VAE if vae_path and vae_path.exists(): print(f"\n[2/5] Converting VAE from {vae_path}") @@ -190,7 +169,7 @@ def convert_full_model( ) else: print(f"\n[2/5] Skipping VAE (not found: {vae_path})") - + # 3. Copy text encoder (LLaMA) if text_encoder_path and text_encoder_path.exists(): print(f"\n[3/5] Copying text encoder from {text_encoder_path}") @@ -198,7 +177,7 @@ def convert_full_model( copy_tokenizer(text_encoder_path, output_dir / "tokenizer", "tokenizer") else: print(f"\n[3/5] Skipping text_encoder (not found: {text_encoder_path})") - + # 4. Copy text encoder 2 (CLIP) if text_encoder_2_path and text_encoder_2_path.exists(): print(f"\n[4/5] Copying text encoder 2 from {text_encoder_2_path}") @@ -206,12 +185,12 @@ def convert_full_model( copy_tokenizer(text_encoder_2_path, output_dir / "tokenizer_2", "tokenizer_2") else: print(f"\n[4/5] Skipping text_encoder_2 (not found: {text_encoder_2_path})") - + # 5. Create model_index.json and scheduler print(f"\n[5/5] Creating model_index.json and scheduler config") create_model_index(output_dir) create_scheduler_config(output_dir) - + print("\n" + "=" * 60) print(f"Conversion complete! Output: {output_dir}") print("\nDirectory structure:") @@ -224,57 +203,27 @@ def convert_full_model( def main(): parser = argparse.ArgumentParser( - description="Convert full Hunyuan-GameCraft-1.0 model to FastVideo diffusers format." - ) - parser.add_argument( - "--input-dir", - type=str, - default=None, - help="Path to Hunyuan-GameCraft-1.0/weights directory" - ) - parser.add_argument( - "--transformer", - type=str, - default=None, - help="Path to transformer checkpoint (mp_rank_00_model_states.pt)" - ) - parser.add_argument( - "--vae", - type=str, - default=None, - help="Path to VAE checkpoint" - ) - parser.add_argument( - "--text-encoder", - type=str, - default=None, - help="Path to text encoder (LLaVA-LLaMA) directory" - ) - parser.add_argument( - "--text-encoder-2", - type=str, - default=None, - help="Path to text encoder 2 (CLIP) directory" - ) - parser.add_argument( - "--output-dir", - type=str, - default="official_weights/hunyuan-gamecraft-diffusers", - help="Output directory for converted model" - ) - parser.add_argument( - "--verbose", - "-v", - action="store_true", - help="Print detailed conversion info" - ) - + description="Convert full Hunyuan-GameCraft-1.0 model to FastVideo diffusers format.") + parser.add_argument("--input-dir", type=str, default=None, help="Path to Hunyuan-GameCraft-1.0/weights directory") + parser.add_argument("--transformer", + type=str, + default=None, + help="Path to transformer checkpoint (mp_rank_00_model_states.pt)") + parser.add_argument("--vae", type=str, default=None, help="Path to VAE checkpoint") + parser.add_argument("--text-encoder", type=str, default=None, help="Path to text encoder (LLaVA-LLaMA) directory") + parser.add_argument("--text-encoder-2", type=str, default=None, help="Path to text encoder 2 (CLIP) directory") + parser.add_argument("--output-dir", + type=str, + default="official_weights/hunyuan-gamecraft-diffusers", + help="Output directory for converted model") + parser.add_argument("--verbose", "-v", action="store_true", help="Print detailed conversion info") + args = parser.parse_args() - + # Require either input-dir or individual paths if args.input_dir is None and args.transformer is None: parser.error("Either --input-dir or individual component paths must be provided") - + convert_full_model( transformer_path=Path(args.transformer) if args.transformer else None, vae_path=Path(args.vae) if args.vae else None, diff --git a/scripts/checkpoint_conversion/convert_gamecraft_vae.py b/scripts/checkpoint_conversion/convert_gamecraft_vae.py index 32982b5e4..de44fb8a6 100644 --- a/scripts/checkpoint_conversion/convert_gamecraft_vae.py +++ b/scripts/checkpoint_conversion/convert_gamecraft_vae.py @@ -57,9 +57,7 @@ def convert_gamecraft_vae( def main(): - parser = argparse.ArgumentParser( - description="Convert GameCraft VAE weights to FastVideo safetensors format." - ) + parser = argparse.ArgumentParser(description="Convert GameCraft VAE weights to FastVideo safetensors format.") parser.add_argument( "--input", type=str, diff --git a/scripts/checkpoint_conversion/convert_gamecraft_weights.py b/scripts/checkpoint_conversion/convert_gamecraft_weights.py index e6dd28281..97a208454 100644 --- a/scripts/checkpoint_conversion/convert_gamecraft_weights.py +++ b/scripts/checkpoint_conversion/convert_gamecraft_weights.py @@ -23,7 +23,6 @@ from pathlib import Path import torch from safetensors.torch import save_file - # Mapping from official GameCraft naming to FastVideo naming # GameCraft weights are already close to FastVideo format, but need minor adjustments PARAM_NAME_MAP: dict[str, str] = { @@ -32,45 +31,46 @@ PARAM_NAME_MAP: dict[str, str] = { r"^(.*)\.img_mlp\.fc2\.(.*)$": r"\1.img_mlp.fc_out.\2", r"^(.*)\.txt_mlp\.fc1\.(.*)$": r"\1.txt_mlp.fc_in.\2", r"^(.*)\.txt_mlp\.fc2\.(.*)$": r"\1.txt_mlp.fc_out.\2", - + # Single block MLP naming r"^single_blocks\.(\d+)\.mlp\.fc1\.(.*)$": r"single_blocks.\1.mlp.fc_in.\2", r"^single_blocks\.(\d+)\.mlp\.fc2\.(.*)$": r"single_blocks.\1.mlp.fc_out.\2", - + # Token refiner naming - rename individual_token_refiner.blocks to refiner_blocks # This is applied first, then subsequent mappings handle the sub-components r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.(.*)$": r"txt_in.refiner_blocks.\1.\2", - + # Vector in naming r"^vector_in\.in_layer\.(.*)$": r"vector_in.fc_in.\1", r"^vector_in\.out_layer\.(.*)$": r"vector_in.fc_out.\1", - + # Time embedder naming r"^time_in\.mlp\.0\.(.*)$": r"time_in.mlp.fc_in.\1", r"^time_in\.mlp\.2\.(.*)$": r"time_in.mlp.fc_out.\1", - + # Guidance embedder naming (if present) r"^guidance_in\.mlp\.0\.(.*)$": r"guidance_in.mlp.fc_in.\1", r"^guidance_in\.mlp\.2\.(.*)$": r"guidance_in.mlp.fc_out.\1", - + # Final layer adaLN modulation r"^final_layer\.adaLN_modulation\.1\.(.*)$": r"final_layer.adaLN_modulation.linear.\1", - + # Text in c_embedder (context embedder) naming r"^txt_in\.c_embedder\.linear_1\.(.*)$": r"txt_in.c_embedder.fc_in.\1", r"^txt_in\.c_embedder\.linear_2\.(.*)$": r"txt_in.c_embedder.fc_out.\1", - + # Text in t_embedder (time embedder) naming r"^txt_in\.t_embedder\.mlp\.0\.(.*)$": r"txt_in.t_embedder.mlp.fc_in.\1", r"^txt_in\.t_embedder\.mlp\.2\.(.*)$": r"txt_in.t_embedder.mlp.fc_out.\1", - + # Refiner block MLP naming - need to match the actual checkpoint names r"^txt_in\.refiner_blocks\.(\d+)\.mlp\.fc1\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_in.\2", r"^txt_in\.refiner_blocks\.(\d+)\.mlp\.fc2\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_out.\2", - + # Refiner block adaLN modulation naming - matches .adaLN_modulation.1.weight/bias - r"^txt_in\.refiner_blocks\.(\d+)\.adaLN_modulation\.1\.(.*)$": r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2", - + r"^txt_in\.refiner_blocks\.(\d+)\.adaLN_modulation\.1\.(.*)$": + r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2", + # Camera net - keep as-is (already correct) # r"^camera_net\.(.*)$": r"camera_net.\1", } @@ -95,7 +95,7 @@ def apply_mapping(key: str) -> str: def load_deepspeed_checkpoint(path: Path) -> dict[str, torch.Tensor]: """Load weights from a DeepSpeed checkpoint file.""" checkpoint = torch.load(path, map_location="cpu") - + if "module" in checkpoint: state_dict = checkpoint["module"] elif "state_dict" in checkpoint: @@ -104,10 +104,10 @@ def load_deepspeed_checkpoint(path: Path) -> dict[str, torch.Tensor]: state_dict = checkpoint["model"] else: state_dict = checkpoint - + if not isinstance(state_dict, dict): raise TypeError(f"Expected dict state_dict, got {type(state_dict)}") - + return state_dict @@ -125,14 +125,14 @@ def convert_weights( """ print(f"Loading checkpoint from {input_path}") state_dict = load_deepspeed_checkpoint(input_path) - + print(f"Found {len(state_dict)} parameters") - + # Convert weights converted_state_dict: OrderedDict[str, torch.Tensor] = OrderedDict() renamed_count = 0 unchanged_count = 0 - + for key, value in state_dict.items(): new_key = apply_mapping(key) if new_key != key: @@ -142,18 +142,18 @@ def convert_weights( else: unchanged_count += 1 converted_state_dict[new_key] = value - + print(f"Renamed {renamed_count} parameters, {unchanged_count} unchanged") - + # Create output directory structure transformer_dir = output_dir / "transformer" transformer_dir.mkdir(parents=True, exist_ok=True) - + # Save as safetensors output_path = transformer_dir / "diffusion_pytorch_model.safetensors" print(f"Saving converted weights to {output_path}") save_file(converted_state_dict, str(output_path)) - + # Save config if requested if save_config: config = { @@ -179,13 +179,13 @@ def convert_weights( with open(config_path, "w") as f: json.dump(config, f, indent=2) print(f"Saved config to {config_path}") - + # Verify camera_net weights camera_keys = [k for k in converted_state_dict.keys() if "camera_net" in k] print(f"\nCamera net weights ({len(camera_keys)}):") for k in camera_keys: print(f" {k}: {converted_state_dict[k].shape}") - + # Verify img_in weights (should be 33 channels) if "img_in.proj.weight" in converted_state_dict: img_in_shape = converted_state_dict["img_in.proj.weight"].shape @@ -196,7 +196,7 @@ def convert_weights( print(f" WARNING: Expected {expected_in_channels} input channels, got {actual_in_channels}") else: print(f" OK: 33 input channels (16 latent + 16 gt_latent + 1 mask)") - + return { "total": len(converted_state_dict), "renamed": renamed_count, @@ -205,48 +205,33 @@ def convert_weights( def main(): - parser = argparse.ArgumentParser( - description="Convert Hunyuan-GameCraft-1.0 weights to FastVideo format." - ) - parser.add_argument( - "--input", - type=str, - required=True, - help="Path to official GameCraft checkpoint (mp_rank_00_model_states.pt)" - ) - parser.add_argument( - "--output", - type=str, - default="official_weights/hunyuan-gamecraft", - help="Output directory for converted weights" - ) - parser.add_argument( - "--no-config", - action="store_true", - help="Don't save config.json" - ) - parser.add_argument( - "--verbose", - "-v", - action="store_true", - help="Print all renamed parameters" - ) - + parser = argparse.ArgumentParser(description="Convert Hunyuan-GameCraft-1.0 weights to FastVideo format.") + parser.add_argument("--input", + type=str, + required=True, + help="Path to official GameCraft checkpoint (mp_rank_00_model_states.pt)") + parser.add_argument("--output", + type=str, + default="official_weights/hunyuan-gamecraft", + help="Output directory for converted weights") + parser.add_argument("--no-config", action="store_true", help="Don't save config.json") + parser.add_argument("--verbose", "-v", action="store_true", help="Print all renamed parameters") + args = parser.parse_args() - + input_path = Path(args.input) output_dir = Path(args.output) - + if not input_path.exists(): raise FileNotFoundError(f"Input file not found: {input_path}") - + stats = convert_weights( input_path=input_path, output_dir=output_dir, save_config=not args.no_config, verbose=args.verbose, ) - + print(f"\nConversion complete!") print(f" Total parameters: {stats['total']}") print(f" Renamed: {stats['renamed']}") diff --git a/scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py b/scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py index 7de24f169..2cc24ca34 100644 --- a/scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py +++ b/scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py @@ -38,7 +38,6 @@ except ImportError: hf_hub_download = None snapshot_download = None - # Parameter name mapping from official GEN3C checkpoint to FastVideo format # Based on fastvideo/configs/models/dits/gen3c.py # The actual GEN3C checkpoint uses AdaLN-LoRA with decomposed projections (index 0 = base, 1 = LoRA) @@ -50,7 +49,7 @@ PARAM_NAMES_MAPPING: dict[str, str] = { r"^net\.t_embedder\.0\.(.*)$": r"time_embed.time_proj.\1", r"^net\.t_embedder\.1\.linear_1\.(.*)$": r"time_embed.t_embedder.linear_1.\1", r"^net\.t_embedder\.1\.linear_2\.(.*)$": r"time_embed.t_embedder.linear_2.\1", - + # Augment sigma embedding (GEN3C-specific) r"^net\.augment_sigma_embedder\.0\.(.*)$": r"augment_sigma_embed.time_proj.\1", r"^net\.augment_sigma_embedder\.1\.linear_1\.(.*)$": r"augment_sigma_embed.t_embedder.linear_1.\1", @@ -66,7 +65,7 @@ PARAM_NAMES_MAPPING: dict[str, str] = { # Transformer blocks: net.blocks.blockN -> transformer_blocks.N # GEN3C uses attn.to_q.0/1 pattern (0=base weight, 1=LoRA weight) - + # Self-attention (block index 0) # Q projection: to_q.0 is linear, to_q.1 is QK norm (RMSNorm applied per-head) r"^net\.blocks\.block(\d+)\.blocks\.0\.block\.attn\.to_q\.0\.(.*)$": r"transformer_blocks.\1.attn1.to_q.\2", @@ -79,7 +78,8 @@ PARAM_NAMES_MAPPING: dict[str, str] = { # Output projection r"^net\.blocks\.block(\d+)\.blocks\.0\.block\.attn\.to_out\.0\.(.*)$": r"transformer_blocks.\1.attn1.to_out.\2", # AdaLN modulation for self-attention - r"^net\.blocks\.block(\d+)\.blocks\.0\.adaLN_modulation\.(.*)$": r"transformer_blocks.\1.adaln_modulation_self_attn.\2", + r"^net\.blocks\.block(\d+)\.blocks\.0\.adaLN_modulation\.(.*)$": + r"transformer_blocks.\1.adaln_modulation_self_attn.\2", # Cross-attention (block index 1) # Q projection: to_q.0 is linear, to_q.1 is QK norm @@ -93,7 +93,8 @@ PARAM_NAMES_MAPPING: dict[str, str] = { # Output projection r"^net\.blocks\.block(\d+)\.blocks\.1\.block\.attn\.to_out\.0\.(.*)$": r"transformer_blocks.\1.attn2.to_out.\2", # AdaLN modulation for cross-attention - r"^net\.blocks\.block(\d+)\.blocks\.1\.adaLN_modulation\.(.*)$": r"transformer_blocks.\1.adaln_modulation_cross_attn.\2", + r"^net\.blocks\.block(\d+)\.blocks\.1\.adaLN_modulation\.(.*)$": + r"transformer_blocks.\1.adaln_modulation_cross_attn.\2", # MLP (block index 2) - simpler naming: layer1, layer2 directly r"^net\.blocks\.block(\d+)\.blocks\.2\.block\.layer1\.(.*)$": r"transformer_blocks.\1.mlp.fc_in.\2", @@ -119,12 +120,12 @@ def apply_mapping(key: str) -> str | None: for pattern in SKIP_PATTERNS: if key.startswith(pattern): return None - + # Apply mapping patterns for pattern, replacement in PARAM_NAMES_MAPPING.items(): if re.match(pattern, key): return re.sub(pattern, replacement, key) - + # If no mapping found, return original key (will be reported) return key @@ -138,7 +139,7 @@ def load_checkpoint(path: Path, mmap: bool = True) -> dict[str, torch.Tensor]: """ # Try memory-mapped loading first (PyTorch 2.1+) load_kwargs: dict = {"map_location": "cpu", "weights_only": False} - + if mmap: # Check if PyTorch version supports mmap (2.1+) try: @@ -152,16 +153,16 @@ def load_checkpoint(path: Path, mmap: bool = True) -> dict[str, torch.Tensor]: print(f" Note: PyTorch {torch.__version__} doesn't support mmap, using standard loading") except (ValueError, IndexError): print(f" Warning: Could not parse PyTorch version {torch.__version__}, skipping mmap") - + checkpoint = torch.load(path, **load_kwargs) - + # Handle different checkpoint formats if isinstance(checkpoint, dict): for key in ("state_dict", "model_state_dict", "model", "ema"): if key in checkpoint: return checkpoint[key] return checkpoint - + return checkpoint @@ -174,13 +175,13 @@ def iterate_checkpoint_keys(path: Path) -> Iterator[str]: # Unfortunately PyTorch doesn't have a great way to do this # So we load the full checkpoint but only keep keys checkpoint = torch.load(path, map_location="meta", weights_only=False) - + if isinstance(checkpoint, dict): for key in ("state_dict", "model_state_dict", "model", "ema"): if key in checkpoint: checkpoint = checkpoint[key] break - + return checkpoint.keys() @@ -189,23 +190,23 @@ def analyze_checkpoint(state_dict: dict[str, torch.Tensor]) -> None: print("\n" + "=" * 80) print("CHECKPOINT ANALYSIS") print("=" * 80) - + # Count parameters total_params = sum(p.numel() for p in state_dict.values()) print(f"\nTotal parameters: {total_params:,} ({total_params / 1e9:.2f}B)") print(f"Total keys: {len(state_dict)}") - + # Analyze key prefixes prefixes: dict[str, int] = {} for key in state_dict: parts = key.split(".") prefix = ".".join(parts[:2]) if len(parts) > 1 else parts[0] prefixes[prefix] = prefixes.get(prefix, 0) + 1 - + print("\nKey prefixes:") for prefix, count in sorted(prefixes.items()): print(f" {prefix}: {count}") - + # Print first 100 keys with shapes print("\nFirst 100 keys with shapes:") for i, (key, value) in enumerate(state_dict.items()): @@ -213,7 +214,7 @@ def analyze_checkpoint(state_dict: dict[str, torch.Tensor]) -> None: print(f" ... and {len(state_dict) - 100} more keys") break print(f" {key}: {list(value.shape)}") - + # Identify GEN3C-specific layers print("\nGEN3C-specific layers:") gen3c_patterns = [ @@ -226,7 +227,7 @@ def analyze_checkpoint(state_dict: dict[str, torch.Tensor]) -> None: if pattern in key: print(f" {key}: {list(state_dict[key].shape)}") break - + print("=" * 80 + "\n") @@ -252,15 +253,15 @@ def convert_weights( converted = OrderedDict() unmapped = [] skipped = [] - + # Get all keys first (so we can delete as we go) keys = list(state_dict.keys()) total = len(keys) - + for i, key in enumerate(keys): value = state_dict[key] new_key = apply_mapping(key) - + if new_key is None: skipped.append(key) if verbose: @@ -281,21 +282,21 @@ def convert_weights( converted[new_key] = value.contiguous() if memory_efficient else value if verbose: print(f" {key} -> {new_key}") - + # Free memory by deleting processed tensor from source if memory_efficient: del state_dict[key] if i % 100 == 0: gc.collect() - + # Progress indicator if (i + 1) % 200 == 0 or i == total - 1: print(f" Processed {i + 1}/{total} tensors...") - + # Final garbage collection if memory_efficient: gc.collect() - + return converted, unmapped, skipped @@ -308,13 +309,13 @@ def write_component( """Write component weights and config to output directory.""" component_dir = output_dir / name component_dir.mkdir(parents=True, exist_ok=True) - + # Save weights output_file = component_dir / "model.safetensors" save_file(weights, str(output_file)) print(f"Saved {name} weights to {output_file}") print(f" {len(weights)} tensors, {sum(t.numel() for t in weights.values()):,} parameters") - + # Save config if config is not None: config_path = component_dir / "config.json" @@ -379,7 +380,7 @@ def download_checkpoint( """Download checkpoint from HuggingFace Hub.""" if hf_hub_download is None: raise RuntimeError("huggingface_hub is required for --download. Install with: uv pip install huggingface_hub") - + print(f"Downloading {filename} from {repo_id}...") path = hf_hub_download( repo_id=repo_id, @@ -401,9 +402,8 @@ def resolve_model_dir( return model_path if snapshot_download is None: - raise RuntimeError( - "huggingface_hub is required to download component source. " - "Install with: uv pip install huggingface_hub") + raise RuntimeError("huggingface_hub is required to download component source. " + "Install with: uv pip install huggingface_hub") print(f"Downloading component source repo: {model_name_or_path}") downloaded = snapshot_download( @@ -451,9 +451,7 @@ def add_inference_components( print(f" Copied {component}: {src} -> {dst}") if missing: - raise FileNotFoundError( - f"Missing required components in source repo {source_dir}: {missing}" - ) + raise FileNotFoundError(f"Missing required components in source repo {source_dir}: {missing}") def patch_gen3c_vae_config(output_dir: Path) -> None: @@ -489,7 +487,7 @@ Examples: python convert_gen3c_to_fastvideo.py --source ./model.pt --output ./gen3c_fastvideo --dtype fp16 """, ) - + parser.add_argument( "--source", type=str, @@ -548,11 +546,9 @@ Examples: "--components-source", type=str, default=None, - help=( - "Optional local path or HF repo id containing diffusers components " - "(vae/text_encoder/tokenizer/scheduler) to copy into output. " - "Example: nvidia/Cosmos-Predict2-2B-Video2World" - ), + help=("Optional local path or HF repo id containing diffusers components " + "(vae/text_encoder/tokenizer/scheduler) to copy into output. " + "Example: nvidia/Cosmos-Predict2-2B-Video2World"), ) parser.add_argument( "--components-cache-dir", @@ -568,21 +564,18 @@ Examples: parser.add_argument( "--components-only", action="store_true", - help=( - "Skip transformer conversion and only add " - "vae/text_encoder/tokenizer/scheduler into --output." - ), + help=("Skip transformer conversion and only add " + "vae/text_encoder/tokenizer/scheduler into --output."), ) - + args = parser.parse_args() - + # Validate arguments if args.components_only: if args.output is None: raise ValueError("--output is required with --components-only") if args.components_source is None: - raise ValueError( - "--components-source is required with --components-only") + raise ValueError("--components-source is required with --components-only") else: if args.download and args.source: raise ValueError("Use either --download or --source, not both") @@ -594,13 +587,10 @@ Examples: if args.components_only: output_dir = Path(args.output) if not output_dir.exists(): - raise FileNotFoundError( - f"--components-only expected existing output directory: {output_dir}" - ) + raise FileNotFoundError(f"--components-only expected existing output directory: {output_dir}") component_source_dir = resolve_model_dir( args.components_source, - cache_dir=Path(args.components_cache_dir) - if args.components_cache_dir else None, + cache_dir=Path(args.components_cache_dir) if args.components_cache_dir else None, ) print("Adding inference components only...") add_inference_components( @@ -611,7 +601,7 @@ Examples: patch_gen3c_vae_config(output_dir) print(f"Done. Components added to {output_dir}") return - + # Parse dtype dtype_map = { "fp32": torch.float32, @@ -619,7 +609,7 @@ Examples: "bf16": torch.bfloat16, } target_dtype = dtype_map[args.dtype] - + # Get checkpoint path if args.download: checkpoint_path = download_checkpoint( @@ -631,63 +621,63 @@ Examples: checkpoint_path = Path(args.source) if not checkpoint_path.exists(): raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") - + # Load checkpoint with memory-mapped loading print(f"Loading checkpoint from {checkpoint_path}...") use_mmap = not args.no_mmap state_dict = load_checkpoint(checkpoint_path, mmap=use_mmap) - + # Analyze if requested if args.analyze: analyze_checkpoint(state_dict) return - + # Analyze before conversion (quick summary only to save memory) print(f"\nCheckpoint has {len(state_dict)} tensors") total_params = sum(p.numel() for p in state_dict.values()) print(f"Total parameters: {total_params:,} ({total_params / 1e9:.2f}B)") - + # Convert weights (memory-efficient mode) print(f"\nConverting weights to {args.dtype}...") converted, unmapped, skipped = convert_weights( - state_dict, + state_dict, verbose=args.verbose, dtype=target_dtype, memory_efficient=True, ) - + # Force garbage collection after conversion del state_dict gc.collect() - + print(f"\nConversion summary:") print(f" Converted: {len(converted)} tensors") print(f" Skipped: {len(skipped)} tensors (dynamic/metadata)") print(f" Unmapped: {len(unmapped)} tensors (kept original names)") - + if unmapped: print("\nWarning: The following keys were not mapped:") for key in unmapped[:20]: print(f" {key}") if len(unmapped) > 20: print(f" ... and {len(unmapped) - 20} more") - + # Write output output_dir = Path(args.output) if output_dir.exists() and not args.force: raise FileExistsError(f"Output directory exists: {output_dir}. Use --force to overwrite.") - + output_dir.mkdir(parents=True, exist_ok=True) - + # Write transformer weights print(f"\nSaving converted weights...") transformer_config = build_transformer_config() write_component(output_dir, "transformer", converted, transformer_config) - + # Free memory after saving del converted gc.collect() - + # Write model_index.json model_index = build_model_index() model_index_path = output_dir / "model_index.json" @@ -700,8 +690,7 @@ Examples: print("\nAdding inference components...") component_source_dir = resolve_model_dir( args.components_source, - cache_dir=Path(args.components_cache_dir) - if args.components_cache_dir else None, + cache_dir=Path(args.components_cache_dir) if args.components_cache_dir else None, ) add_inference_components( source_dir=component_source_dir, @@ -713,7 +702,7 @@ Examples: print("\nNote: Only transformer/model_index were written.") print("FastVideo local loading also requires: vae/, text_encoder/, tokenizer/, scheduler/.") print("Re-run with --components-source to add them automatically.") - + print(f"\nConversion complete! Output saved to {output_dir}") print("\nTo use with FastVideo:") print(f" from fastvideo.models.dits.gen3c import Gen3CTransformer3DModel") diff --git a/scripts/checkpoint_conversion/convert_lingbotworld2_causal_fast.py b/scripts/checkpoint_conversion/convert_lingbotworld2_causal_fast.py index 8a232cffe..4d9f5a768 100644 --- a/scripts/checkpoint_conversion/convert_lingbotworld2_causal_fast.py +++ b/scripts/checkpoint_conversion/convert_lingbotworld2_causal_fast.py @@ -8,7 +8,6 @@ import json import os from pathlib import Path - LATENTS_MEAN = [ -0.7571, -0.7089, @@ -150,9 +149,11 @@ def build_bundle(source_dir: Path, output_dir: Path) -> None: "final_sigmas_type": "zero", }, ) - _write_json(output_dir / "scheduler" / "config.json", json.loads((output_dir / "scheduler" / "scheduler_config.json").read_text())) + _write_json(output_dir / "scheduler" / "config.json", + json.loads((output_dir / "scheduler" / "scheduler_config.json").read_text())) - _symlink_or_replace(source_dir / "models_t5_umt5-xxl-enc-bf16.pth", output_dir / "text_encoder" / "pytorch_model.pt") + _symlink_or_replace(source_dir / "models_t5_umt5-xxl-enc-bf16.pth", + output_dir / "text_encoder" / "pytorch_model.pt") _symlink_or_replace(source_dir / "google" / "umt5-xxl", output_dir / "tokenizer") for shard in sorted((source_dir / "transformers").glob("*.safetensors")): _symlink_or_replace(shard, output_dir / "transformer" / shard.name) diff --git a/scripts/checkpoint_conversion/convert_ltx2_weights.py b/scripts/checkpoint_conversion/convert_ltx2_weights.py index 7cb66cf33..c731ed866 100644 --- a/scripts/checkpoint_conversion/convert_ltx2_weights.py +++ b/scripts/checkpoint_conversion/convert_ltx2_weights.py @@ -36,16 +36,15 @@ try: except ImportError: # pragma: no cover - optional dependency snapshot_download = None - PARAM_NAME_MAP: dict[str, str] = { r"^model\.diffusion_model\.(.*)$": r"\1", } COMPONENT_PREFIXES: dict[str, tuple[str, ...]] = { - "transformer": ("model.diffusion_model.",), - "vae": ("vae.",), - "audio_vae": ("audio_vae.",), - "vocoder": ("vocoder.",), + "transformer": ("model.diffusion_model.", ), + "vae": ("vae.", ), + "audio_vae": ("audio_vae.", ), + "vocoder": ("vocoder.", ), "text_embedding_projection": ("text_embedding_projection.", "model.text_embedding_projection."), } @@ -117,9 +116,7 @@ def _filter_transformer_config(config: dict) -> dict: return filtered -def _build_text_embedding_projection_config( - gemma_model_path: str = "", -) -> dict: +def _build_text_embedding_projection_config(gemma_model_path: str = "", ) -> dict: return { "architectures": ["LTX2GemmaTextEncoderModel"], "hidden_size": 3840, @@ -278,9 +275,7 @@ def convert_components( metadata_config.get("vocoder"), class_name="LTX2Vocoder", ), - "text_embedding_projection": _build_text_embedding_projection_config( - gemma_model_path=gemma_model_path - ), + "text_embedding_projection": _build_text_embedding_projection_config(gemma_model_path=gemma_model_path), } output_dir.mkdir(parents=True, exist_ok=True) @@ -308,9 +303,7 @@ def convert_components( if not required_for_index.issubset(split_weights.keys()): print("Skipping model_index.json; missing diffusers components in weights.") return - vae_class_name = (component_configs.get("vae") or {}).get( - "_class_name", "CausalVideoAutoencoder" - ) + vae_class_name = (component_configs.get("vae") or {}).get("_class_name", "CausalVideoAutoencoder") model_index = _build_model_index( transformer_class_name=transformer_class_name, vae_class_name=vae_class_name, @@ -369,9 +362,7 @@ def copy_gemma_tokenizer(gemma_src: Path, tokenizer_dest: Path) -> None: shutil.copy2(src_path, tokenizer_dest / file_name) copied += 1 if copied == 0: - raise FileNotFoundError( - f"No tokenizer files found in {gemma_src}. Expected at least one tokenizer file." - ) + raise FileNotFoundError(f"No tokenizer files found in {gemma_src}. Expected at least one tokenizer file.") print(f"Copied {copied} tokenizer files to {tokenizer_dest}") @@ -411,10 +402,8 @@ def main() -> None: "--components", type=str, default="", - help=( - "Comma-separated component list to write " - "(transformer,vae,audio_vae,vocoder,text_embedding_projection)." - ), + help=("Comma-separated component list to write " + "(transformer,vae,audio_vae,vocoder,text_embedding_projection)."), ) parser.add_argument( "--gemma-path", @@ -444,11 +433,7 @@ def main() -> None: if args.transformer_only: components_to_write = {"transformer"} elif args.components: - components_to_write = { - component.strip() - for component in args.components.split(",") - if component.strip() - } + components_to_write = {component.strip() for component in args.components.split(",") if component.strip()} gemma_model_path = "" if args.gemma_path: diff --git a/scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py b/scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py index 888781cc2..6e131affd 100644 --- a/scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py +++ b/scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py @@ -75,7 +75,6 @@ import torch from huggingface_hub import hf_hub_download, snapshot_download from safetensors.torch import load_file, save_file - # MagiHuman base arch — keys must be valid `MagiHumanArchConfig` fields. # FastVideo's `TransformerLoader.load` calls `ArchConfig.update_model_arch` # with this dict (minus `_class_name`, `_diffusers_version`) and rejects @@ -91,9 +90,9 @@ MAGI_HUMAN_BASE_ARCH: dict = { "head_dim": 128, "num_query_groups": 8, # Modality channels. - "video_in_channels": 192, # 48 (VAE z_dim) * patch_size product 1*2*2 + "video_in_channels": 192, # 48 (VAE z_dim) * patch_size product 1*2*2 "audio_in_channels": 64, - "text_in_channels": 3584, # T5Gemma-9B encoder hidden size + "text_in_channels": 3584, # T5Gemma-9B encoder hidden size # Block-level switches. "mm_layers": [0, 1, 2, 3, 36, 37, 38, 39], "local_attn_layers": [], @@ -110,7 +109,6 @@ MAGI_HUMAN_BASE_ARCH: dict = { "tread_end_layer_idx": 25, } - SCHEDULER_CONFIG: dict = { "_class_name": "FlowUniPCMultistepScheduler", "_diffusers_version": "0.33.0", @@ -125,7 +123,6 @@ SCHEDULER_CONFIG: dict = { "flow_shift": 5.0, } - MAX_SHARD_BYTES = 5 * 1024 * 1024 * 1024 # 5 GB shards, matches HF defaults @@ -316,10 +313,7 @@ def _write_model_index( ) -> None: pipeline_class = "MagiHumanPipeline" if include_sr_transformer: - pipeline_class = ( - "MagiHumanSR1080pPipeline" - if sr_subfolder == "1080p_sr" else "MagiHumanSRPipeline" - ) + pipeline_class = ("MagiHumanSR1080pPipeline" if sr_subfolder == "1080p_sr" else "MagiHumanSRPipeline") index = { "_class_name": pipeline_class, "_diffusers_version": "0.33.0", @@ -374,13 +368,11 @@ def _bundle_sa_audio_vae(out_dir: Path, source_repo: str = "stabilityai/stable-a https://huggingface.co/stabilityai/stable-audio-open-1.0. """ print(f" fetching audio VAE from {source_repo} (gated) ...") - token = ( - os.environ.get("HF_TOKEN") - or os.environ.get("HUGGINGFACE_HUB_TOKEN") - or os.environ.get("HF_API_KEY") - ) + token = (os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get("HF_API_KEY")) local = snapshot_download( - repo_id=source_repo, token=token, allow_patterns=["vae/*"], + repo_id=source_repo, + token=token, + allow_patterns=["vae/*"], ) src = Path(local) / "vae" if not src.exists(): @@ -399,11 +391,7 @@ def _bundle_text_encoder(out_dir: Path, source_repo: str = "google/t5gemma-9b-9b accepted terms of use for the repo. """ print(f" fetching text encoder from {source_repo} (gated) ...") - token = ( - os.environ.get("HF_TOKEN") - or os.environ.get("HUGGINGFACE_HUB_TOKEN") - or os.environ.get("HF_API_KEY") - ) + token = (os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get("HF_API_KEY")) local = snapshot_download( repo_id=source_repo, token=token, @@ -428,8 +416,9 @@ def _bundle_text_encoder(out_dir: Path, source_repo: str = "google/t5gemma-9b-9b dst_tokenizer.mkdir(parents=True, exist_ok=True) for fname in src.iterdir(): - if fname.name in {"tokenizer.model", "tokenizer.json", "tokenizer_config.json", - "special_tokens_map.json", "spiece.model"}: + if fname.name in { + "tokenizer.model", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "spiece.model" + }: shutil.copy(fname, dst_tokenizer / fname.name) else: shutil.copy(fname, dst_encoder / fname.name) @@ -462,29 +451,28 @@ def main() -> None: parser.add_argument( "--cast-bf16", action="store_true", - help=( - "Cast fp32 DiT weights to bfloat16 on save. Recommended for the " - "distill subfolder (61 GB fp32 upstream -> 30 GB bf16 artifact). " - "Keeps norms, RoPE bands, and other precision-sensitive tensors " - "in fp32." - ), + help=("Cast fp32 DiT weights to bfloat16 on save. Recommended for the " + "distill subfolder (61 GB fp32 upstream -> 30 GB bf16 artifact). " + "Keeps norms, RoPE bands, and other precision-sensitive tensors " + "in fp32."), ) parser.add_argument( "--bundle-text-encoder", action="store_true", help="Download google/t5gemma-9b-9b-ul2 into /text_encoder/ and tokenizer/. " - "Requires a write-scoped HF token with accepted terms of use.", + "Requires a write-scoped HF token with accepted terms of use.", ) parser.add_argument( "--bundle-audio-vae", action="store_true", help="Download stabilityai/stable-audio-open-1.0 VAE into /audio_vae/. " - "Requires HF terms accepted for the Stability AI gated repo.", + "Requires HF terms accepted for the Stability AI gated repo.", ) parser.add_argument( "--sr-source", default=None, - help="Optional HF repo id or local directory containing SR DiT shards. When set, writes /sr_transformer/.", + help= + "Optional HF repo id or local directory containing SR DiT shards. When set, writes /sr_transformer/.", ) parser.add_argument( "--sr-subfolder", diff --git a/scripts/checkpoint_conversion/convert_turbodiffusion_i2v_to_diffusers.py b/scripts/checkpoint_conversion/convert_turbodiffusion_i2v_to_diffusers.py index 3ed230b82..ccc997457 100644 --- a/scripts/checkpoint_conversion/convert_turbodiffusion_i2v_to_diffusers.py +++ b/scripts/checkpoint_conversion/convert_turbodiffusion_i2v_to_diffusers.py @@ -92,21 +92,20 @@ def reshape_patch_embedding(tensor: torch.Tensor, target_shape: tuple) -> torch. def get_reference_shapes(reference_repo: str) -> dict: print(f"Downloading reference model shapes from {reference_repo}...") - + local_dir = snapshot_download( repo_id=reference_repo, allow_patterns=["transformer/config.json", "transformer/diffusion_pytorch_model*.safetensors"], - local_dir_use_symlinks=False - ) - + local_dir_use_symlinks=False) + weight_files = glob.glob(os.path.join(local_dir, "transformer", "*.safetensors")) shapes = {} - + for wf in weight_files: with safe_open(wf, framework="pt") as f: for key in f.keys(): shapes[key] = f.get_tensor(key).shape - + return shapes, local_dir @@ -115,73 +114,78 @@ def convert_checkpoint(input_path: str, output_dir: str, ref_shapes: dict, model print(f"\n{'='*60}") print(f"Converting {model_name}: {input_path}") print(f"{'='*60}") - + turbo_state_dict = torch.load(input_path, map_location="cpu", weights_only=True) print(f"Loaded {len(turbo_state_dict)} keys") - + converted_state_dict = {} skipped_keys = [] - + for turbo_key, tensor in turbo_state_dict.items(): if should_skip_key(turbo_key): skipped_keys.append(turbo_key) continue - + new_key = convert_key(turbo_key) - + if "patch_embedding" in new_key and new_key in ref_shapes: target_shape = ref_shapes[new_key] if tensor.shape != target_shape: print(f"Reshaping {new_key}: {tensor.shape} -> {target_shape}") tensor = reshape_patch_embedding(tensor, target_shape) - + if new_key in ref_shapes: if tensor.shape != ref_shapes[new_key]: print(f"WARNING: Shape mismatch for {new_key}: got {tensor.shape}, expected {ref_shapes[new_key]}") - + converted_state_dict[new_key] = tensor - + print(f"Converted: {len(converted_state_dict)} keys, Skipped: {len(skipped_keys)} keys") - + os.makedirs(output_dir, exist_ok=True) output_path = os.path.join(output_dir, "diffusion_pytorch_model.safetensors") print(f"Saving to {output_path}...") save_file(converted_state_dict, output_path) - + return converted_state_dict def main(): parser = argparse.ArgumentParser(description="Convert TurboDiffusion I2V checkpoints to Diffusers format") - parser.add_argument("--high_noise_path", type=str, required=True, + parser.add_argument("--high_noise_path", + type=str, + required=True, help="Path to high-noise TurboDiffusion .pth checkpoint") - parser.add_argument("--low_noise_path", type=str, required=True, + parser.add_argument("--low_noise_path", + type=str, + required=True, help="Path to low-noise TurboDiffusion .pth checkpoint") - parser.add_argument("--output_dir", type=str, required=True, - help="Output directory for converted safetensors") - parser.add_argument("--reference_repo", type=str, default="Wan-AI/Wan2.2-I2V-A14B-Diffusers", + parser.add_argument("--output_dir", type=str, required=True, help="Output directory for converted safetensors") + parser.add_argument("--reference_repo", + type=str, + default="Wan-AI/Wan2.2-I2V-A14B-Diffusers", help="Reference HF repo to get expected tensor shapes") - + args = parser.parse_args() - + # Get reference shapes ref_shapes, ref_local_dir = get_reference_shapes(args.reference_repo) print(f"Got {len(ref_shapes)} reference shapes") - + # Convert high-noise model high_noise_output = os.path.join(args.output_dir, "transformer_high") convert_checkpoint(args.high_noise_path, high_noise_output, ref_shapes, "high-noise") - + # Convert low-noise model low_noise_output = os.path.join(args.output_dir, "transformer_low") convert_checkpoint(args.low_noise_path, low_noise_output, ref_shapes, "low-noise") - + # Copy config.json to both src_config = os.path.join(ref_local_dir, "transformer", "config.json") shutil.copy(src_config, os.path.join(high_noise_output, "config.json")) shutil.copy(src_config, os.path.join(low_noise_output, "config.json")) print("Copied config.json to both transformer directories") - + print(f"\n{'='*60}") print(f"Conversion complete!") print(f"High-noise model: {high_noise_output}") diff --git a/scripts/checkpoint_conversion/convert_turbodiffusion_to_diffusers.py b/scripts/checkpoint_conversion/convert_turbodiffusion_to_diffusers.py index 077add50e..72f405a29 100644 --- a/scripts/checkpoint_conversion/convert_turbodiffusion_to_diffusers.py +++ b/scripts/checkpoint_conversion/convert_turbodiffusion_to_diffusers.py @@ -94,103 +94,102 @@ def reshape_patch_embedding(tensor: torch.Tensor, target_shape: tuple) -> torch. def get_reference_shapes(reference_repo: str) -> dict: """Download reference model and get expected shapes for patch_embedding.""" print(f"Downloading reference model shapes from {reference_repo}...") - + # Download just the transformer config and a weight file to get shapes local_dir = snapshot_download( repo_id=reference_repo, allow_patterns=["transformer/config.json", "transformer/diffusion_pytorch_model*.safetensors"], - local_dir_use_symlinks=False - ) - - # Load the first safetensors file to get shapes + local_dir_use_symlinks=False) + + # Load the first safetensors file to get shapes weight_files = glob.glob(os.path.join(local_dir, "transformer", "*.safetensors")) shapes = {} - + for wf in weight_files: with safe_open(wf, framework="pt") as f: for key in f.keys(): shapes[key] = f.get_tensor(key).shape - + return shapes def main(): parser = argparse.ArgumentParser(description="Convert TurboDiffusion checkpoint to Diffusers format") - parser.add_argument("--input_path", type=str, required=True, - help="Path to TurboDiffusion .pth checkpoint") - parser.add_argument("--output_dir", type=str, required=True, - help="Output directory for converted safetensors") - parser.add_argument("--reference_repo", type=str, default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + parser.add_argument("--input_path", type=str, required=True, help="Path to TurboDiffusion .pth checkpoint") + parser.add_argument("--output_dir", type=str, required=True, help="Output directory for converted safetensors") + parser.add_argument("--reference_repo", + type=str, + default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", help="Reference HF repo to get expected tensor shapes") - parser.add_argument("--skip_sla_weights", action="store_true", default=False, + parser.add_argument("--skip_sla_weights", + action="store_true", + default=False, help="Skip SLA-specific weights (proj_l) that aren't in base model") - + args = parser.parse_args() - + # Load TurboDiffusion checkpoint print(f"Loading TurboDiffusion checkpoint from {args.input_path}...") turbo_state_dict = torch.load(args.input_path, map_location="cpu", weights_only=True) print(f"Loaded {len(turbo_state_dict)} keys") - + # Get reference shapes for reshaping ref_shapes = get_reference_shapes(args.reference_repo) print(f"Got {len(ref_shapes)} reference shapes") - + # Convert keys and reshape tensors converted_state_dict = {} skipped_keys = [] - + for turbo_key, tensor in turbo_state_dict.items(): # Skip SLA-specific weights if requested if args.skip_sla_weights and should_skip_key(turbo_key): skipped_keys.append(turbo_key) continue - + # Convert key name new_key = convert_key(turbo_key) - + # Reshape patch_embedding if needed if "patch_embedding" in new_key and new_key in ref_shapes: target_shape = ref_shapes[new_key] if tensor.shape != target_shape: print(f"Reshaping {new_key}: {tensor.shape} -> {target_shape}") tensor = reshape_patch_embedding(tensor, target_shape) - + # Verify shape matches reference if available if new_key in ref_shapes: if tensor.shape != ref_shapes[new_key]: print(f"WARNING: Shape mismatch for {new_key}: got {tensor.shape}, expected {ref_shapes[new_key]}") - + converted_state_dict[new_key] = tensor - + print(f"\nConversion summary:") print(f" Converted: {len(converted_state_dict)} keys") print(f" Skipped (SLA): {len(skipped_keys)} keys") - + if skipped_keys: print(f"\nSkipped SLA keys (first 5):") for k in skipped_keys[:5]: print(f" - {k}") - + # Create output directory os.makedirs(args.output_dir, exist_ok=True) - + # Save as safetensors output_path = os.path.join(args.output_dir, "diffusion_pytorch_model.safetensors") print(f"\nSaving to {output_path}...") save_file(converted_state_dict, output_path) - + # Copy config.json from reference - ref_local = snapshot_download( - repo_id=args.reference_repo, - allow_patterns=["transformer/config.json"], - local_dir_use_symlinks=False - ) + ref_local = snapshot_download(repo_id=args.reference_repo, + allow_patterns=["transformer/config.json"], + local_dir_use_symlinks=False) src_config = os.path.join(ref_local, "transformer", "config.json") dst_config = os.path.join(args.output_dir, "config.json") shutil.copy(src_config, dst_config) print(f"Copied config.json") - + print(f"\nDone! Converted weights saved to: {args.output_dir}") print(f"\nNext steps:") print(f" 1. Use create_hf_repo.py to create a complete diffusers repo:") diff --git a/scripts/checkpoint_conversion/create_hf_repo.py b/scripts/checkpoint_conversion/create_hf_repo.py index 097dbb7dc..d7ecdc78d 100644 --- a/scripts/checkpoint_conversion/create_hf_repo.py +++ b/scripts/checkpoint_conversion/create_hf_repo.py @@ -4,60 +4,50 @@ import shutil import glob from huggingface_hub import snapshot_download, upload_folder, create_repo + def main(): parser = argparse.ArgumentParser(description="Download a HF Diffusers repo and replace its transformer weights.") parser.add_argument( - "--repo_id", - type=str, - required=True, - help="The Hugging Face repository ID to download (e.g., 'stabilityai/stable-diffusion-xl-base-1.0')." - ) + "--repo_id", + type=str, + required=True, + help="The Hugging Face repository ID to download (e.g., 'stabilityai/stable-diffusion-xl-base-1.0').") + parser.add_argument("--local_dir", + type=str, + required=True, + help="The local directory where the repository will be downloaded and modified.") + parser.add_argument("--checkpoint_dir", + type=str, + required=True, + help="The directory containing the new transformer model weights (checkpoints) to inject.") parser.add_argument( - "--local_dir", - type=str, - required=True, - help="The local directory where the repository will be downloaded and modified." - ) - parser.add_argument( - "--checkpoint_dir", - type=str, - required=True, - help="The directory containing the new transformer model weights (checkpoints) to inject." - ) - parser.add_argument( - "--component_name", - type=str, - default="transformer", + "--component_name", + type=str, + default="transformer", help="The name of the component subfolder in the diffusers repo to replace weights for (default: 'transformer')." ) parser.add_argument( "--weight_file_name", type=str, default="diffusion_pytorch_model.safetensors", - help="The target weight filename in the component directory (default: 'diffusion_pytorch_model.safetensors')." - ) - parser.add_argument( - "--ignore_patterns", - nargs="+", - default=[], - help="Patterns to ignore when downloading the repo (passed to snapshot_download)." - ) - parser.add_argument( - "--push_to_hub", - action="store_true", - help="If set, upload the modified local directory back to the Hugging Face Hub (to --repo_id)." - ) - parser.add_argument( - "--hub_commit_message", - type=str, - default="Update transformer weights", - help="Commit message for the Hub upload." - ) + help="The target weight filename in the component directory (default: 'diffusion_pytorch_model.safetensors').") + parser.add_argument("--ignore_patterns", + nargs="+", + default=[], + help="Patterns to ignore when downloading the repo (passed to snapshot_download).") + parser.add_argument("--push_to_hub", + action="store_true", + help="If set, upload the modified local directory back to the Hugging Face Hub (to --repo_id).") + parser.add_argument("--hub_commit_message", + type=str, + default="Update transformer weights", + help="Commit message for the Hub upload.") parser.add_argument( "--upload_repo_id", type=str, default=None, - help="The Hugging Face repository ID to upload to. If not provided, defaults to --repo_id. Creates the repo if it doesn't exist." + help= + "The Hugging Face repository ID to upload to. If not provided, defaults to --repo_id. Creates the repo if it doesn't exist." ) args = parser.parse_args() @@ -81,7 +71,7 @@ def main(): # Find weight files in checkpoint_dir source_weights = glob.glob(os.path.join(args.checkpoint_dir, "*.safetensors")) + \ glob.glob(os.path.join(args.checkpoint_dir, "*.bin")) - + if not source_weights: print(f"Error: No .safetensors or .bin weights found in '{args.checkpoint_dir}'.") return @@ -92,7 +82,7 @@ def main(): existing_weights = glob.glob(os.path.join(target_component_dir, "*.safetensors")) + \ glob.glob(os.path.join(target_component_dir, "*.bin")) + \ glob.glob(os.path.join(target_component_dir, "*.index.json")) # Remove index if present - + for ew in existing_weights: print(f"Removing existing weight file: {ew}") os.remove(ew) @@ -101,12 +91,12 @@ def main(): for src in source_weights: filename = os.path.basename(src) dst = os.path.join(target_component_dir, filename) - + # If there is only one file and it doesn't match the standard name, maybe we should rename it? if len(source_weights) == 1 and filename != args.weight_file_name: - print(f"Renaming '{filename}' to '{args.weight_file_name}'...") - dst = os.path.join(target_component_dir, args.weight_file_name) - + print(f"Renaming '{filename}' to '{args.weight_file_name}'...") + dst = os.path.join(target_component_dir, args.weight_file_name) + print(f"Copying '{src}' to '{dst}'...") shutil.copy2(src, dst) @@ -116,19 +106,16 @@ def main(): if args.push_to_hub: target_repo_id = args.upload_repo_id if args.upload_repo_id else args.repo_id print(f"Uploading to Hugging Face Hub: {target_repo_id}...") - + try: # Ensure the repo exists create_repo(target_repo_id, exist_ok=True) - - upload_folder( - repo_id=target_repo_id, - folder_path=args.local_dir, - commit_message=args.hub_commit_message - ) + + upload_folder(repo_id=target_repo_id, folder_path=args.local_dir, commit_message=args.hub_commit_message) print("Upload complete!") except Exception as e: print(f"Error uploading to Hub: {e}") + if __name__ == "__main__": main() diff --git a/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py b/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py index cc83e1ddc..449f51afe 100644 --- a/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py +++ b/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py @@ -28,7 +28,6 @@ from huggingface_hub import save_torch_state_dict from safetensors import safe_open from safetensors.torch import load_file - OFFICIAL_TO_DIFFUSERS_MAPPING: dict[str, str] = { r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.linear_1.\1", r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.linear_2.\1", @@ -194,8 +193,14 @@ def analyze(source: Path) -> None: def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--source", type=Path, required=True, help="DreamX raw transformer directory or safetensors file") - parser.add_argument("--output", type=Path, required=True, help="Output model root; transformer/ is created inside it") + parser.add_argument("--source", + type=Path, + required=True, + help="DreamX raw transformer directory or safetensors file") + parser.add_argument("--output", + type=Path, + required=True, + help="Output model root; transformer/ is created inside it") parser.add_argument("--max-shard-size", default="10GB") parser.add_argument( "--component-source", diff --git a/scripts/checkpoint_conversion/extract_llava_text_encoder.py b/scripts/checkpoint_conversion/extract_llava_text_encoder.py index 11f616ce9..b202aec4c 100644 --- a/scripts/checkpoint_conversion/extract_llava_text_encoder.py +++ b/scripts/checkpoint_conversion/extract_llava_text_encoder.py @@ -30,17 +30,17 @@ def extract_text_encoder( output_dir: Path, ) -> dict: """Extract LlamaModel from LLaVA weights.""" - + # Find all safetensor shards shard_files = sorted(input_dir.glob("model-*.safetensors")) if not shard_files: raise FileNotFoundError(f"No safetensor files found in {input_dir}") - + print(f"Found {len(shard_files)} weight shards") - + # Collect all weights, filtering for language_model.model prefix extracted_weights: OrderedDict[str, torch.Tensor] = OrderedDict() - + for shard_file in shard_files: print(f"Processing {shard_file.name}...") with safe_open(str(shard_file), framework="pt") as f: @@ -50,17 +50,17 @@ def extract_text_encoder( # Strip the prefix new_key = key.replace("language_model.model.", "") extracted_weights[new_key] = f.get_tensor(key) - + print(f"Extracted {len(extracted_weights)} weights") - + # Save to output directory output_dir.mkdir(parents=True, exist_ok=True) - + # Save as single file (it should be <10GB which is fine for safetensors) output_file = output_dir / "model.safetensors" print(f"Saving to {output_file}...") save_file(extracted_weights, str(output_file)) - + # Create config.json for LlamaModel config = { "_name_or_path": "llava-llama-3-8b-v1_1-text_encoder", @@ -90,15 +90,15 @@ def extract_text_encoder( "use_cache": True, "vocab_size": 128320 } - + with open(output_dir / "config.json", "w") as f: json.dump(config, f, indent=2) print(f"Saved config.json") - + # Copy tokenizer files if they exist tokenizer_files = [ "tokenizer.json", - "tokenizer_config.json", + "tokenizer_config.json", "special_tokens_map.json", ] for fname in tokenizer_files: @@ -107,7 +107,7 @@ def extract_text_encoder( import shutil shutil.copy2(src, output_dir / fname) print(f"Copied {fname}") - + return {"total": len(extracted_weights)} @@ -116,45 +116,45 @@ def extract_clip_text_encoder( output_dir: Path, ) -> dict: """Extract CLIPTextModel from full CLIP model weights.""" - + # Find weight file weight_files = list(input_dir.glob("model.safetensors")) + \ list(input_dir.glob("pytorch_model.bin")) - + if not weight_files: raise FileNotFoundError(f"No weight files found in {input_dir}") - + weight_file = weight_files[0] print(f"Loading {weight_file}...") - + if weight_file.suffix == ".safetensors": from safetensors.torch import load_file state_dict = load_file(str(weight_file)) else: state_dict = torch.load(weight_file, map_location="cpu") - + # Extract only text_model weights extracted_weights: OrderedDict[str, torch.Tensor] = OrderedDict() - + for key, value in state_dict.items(): # CLIP full model has text_model.XXX, we want just XXX if key.startswith("text_model."): new_key = key.replace("text_model.", "") extracted_weights[new_key] = value - + if not extracted_weights: # Maybe it's already just the text model extracted_weights = OrderedDict(state_dict) - + print(f"Extracted {len(extracted_weights)} weights") - - # Save to output directory + + # Save to output directory output_dir.mkdir(parents=True, exist_ok=True) - + output_file = output_dir / "model.safetensors" print(f"Saving to {output_file}...") save_file(extracted_weights, str(output_file)) - + # Create config.json for CLIPTextModel config = { "_name_or_path": "openai/clip-vit-large-patch14", @@ -179,11 +179,11 @@ def extract_clip_text_encoder( "transformers_version": "4.48.0", "vocab_size": 49408 } - + with open(output_dir / "config.json", "w") as f: json.dump(config, f, indent=2) print(f"Saved config.json") - + # Copy tokenizer files if they exist tokenizer_files = [ "tokenizer.json", @@ -198,47 +198,33 @@ def extract_clip_text_encoder( import shutil shutil.copy2(src, output_dir / fname) print(f"Copied {fname}") - + return {"total": len(extracted_weights)} def main(): - parser = argparse.ArgumentParser( - description="Extract text encoder from LLaVA or CLIP model weights." - ) - parser.add_argument( - "--input", - type=str, - required=True, - help="Input directory containing model weights" - ) - parser.add_argument( - "--output", - type=str, - required=True, - help="Output directory for extracted weights" - ) - parser.add_argument( - "--type", - type=str, - choices=["llama", "clip"], - default="llama", - help="Type of text encoder to extract" - ) - + parser = argparse.ArgumentParser(description="Extract text encoder from LLaVA or CLIP model weights.") + parser.add_argument("--input", type=str, required=True, help="Input directory containing model weights") + parser.add_argument("--output", type=str, required=True, help="Output directory for extracted weights") + parser.add_argument("--type", + type=str, + choices=["llama", "clip"], + default="llama", + help="Type of text encoder to extract") + args = parser.parse_args() - + input_dir = Path(args.input) output_dir = Path(args.output) - + if not input_dir.exists(): raise FileNotFoundError(f"Input directory not found: {input_dir}") - + if args.type == "llama": stats = extract_text_encoder(input_dir, output_dir) else: stats = extract_clip_text_encoder(input_dir, output_dir) - + print(f"\nExtraction complete! Total weights: {stats['total']}") diff --git a/scripts/checkpoint_conversion/longcat_to_fastvideo.py b/scripts/checkpoint_conversion/longcat_to_fastvideo.py index 7a9996724..19a80601c 100644 --- a/scripts/checkpoint_conversion/longcat_to_fastvideo.py +++ b/scripts/checkpoint_conversion/longcat_to_fastvideo.py @@ -35,12 +35,12 @@ def split_qkv(qkv_weight: torch.Tensor, qkv_bias: torch.Tensor | None = None): """Split fused QKV projection into separate Q, K, V.""" dim = qkv_weight.shape[0] // 3 q, k, v = torch.chunk(qkv_weight, 3, dim=0) - + if qkv_bias is not None: q_bias, k_bias, v_bias = torch.chunk(qkv_bias, 3, dim=0) else: q_bias = k_bias = v_bias = None - + return (q, k, v), (q_bias, k_bias, v_bias) @@ -48,12 +48,12 @@ def split_kv(kv_weight: torch.Tensor, kv_bias: torch.Tensor | None = None): """Split fused KV projection into separate K, V.""" dim = kv_weight.shape[0] // 2 k, v = torch.chunk(kv_weight, 2, dim=0) - + if kv_bias is not None: k_bias, v_bias = torch.chunk(kv_bias, 2, dim=0) else: k_bias = v_bias = None - + return (k, v), (k_bias, v_bias) @@ -68,171 +68,171 @@ def convert_transformer_weights(source_weights: dict[str, torch.Tensor]) -> dict """ converted = OrderedDict() processed_keys = set() - + print(" Converting transformer weights...") - + for key, value in tqdm(source_weights.items(), desc=" Processing parameters"): if key in processed_keys: continue - + # === Embedders === if key.startswith("x_embedder."): new_key = key.replace("x_embedder.", "patch_embed.") converted[new_key] = value - + elif key.startswith("t_embedder.mlp.0."): new_key = key.replace("t_embedder.mlp.0.", "time_embedder.linear_1.") converted[new_key] = value - + elif key.startswith("t_embedder.mlp.2."): new_key = key.replace("t_embedder.mlp.2.", "time_embedder.linear_2.") converted[new_key] = value - + elif key.startswith("y_embedder.y_proj.0."): new_key = key.replace("y_embedder.y_proj.0.", "caption_embedder.linear_1.") converted[new_key] = value - + elif key.startswith("y_embedder.y_proj.2."): new_key = key.replace("y_embedder.y_proj.2.", "caption_embedder.linear_2.") converted[new_key] = value - + # === Self-Attention QKV Splitting === elif ".attn.qkv.weight" in key: block_idx = key.split(".")[1] qkv_weight = value qkv_bias_key = key.replace(".weight", ".bias") qkv_bias = source_weights.get(qkv_bias_key) - + (q, k, v), (q_bias, k_bias, v_bias) = split_qkv(qkv_weight, qkv_bias) - + converted[f"blocks.{block_idx}.self_attn.to_q.weight"] = q converted[f"blocks.{block_idx}.self_attn.to_k.weight"] = k converted[f"blocks.{block_idx}.self_attn.to_v.weight"] = v - + if q_bias is not None: converted[f"blocks.{block_idx}.self_attn.to_q.bias"] = q_bias converted[f"blocks.{block_idx}.self_attn.to_k.bias"] = k_bias converted[f"blocks.{block_idx}.self_attn.to_v.bias"] = v_bias - + processed_keys.add(key) if qkv_bias is not None: processed_keys.add(qkv_bias_key) - + elif ".attn.qkv.bias" in key: continue - + elif ".attn.proj." in key: new_key = key.replace(".attn.proj.", ".self_attn.to_out.") converted[new_key] = value - + elif ".attn.q_norm." in key or ".attn.k_norm." in key: new_key = key.replace(".attn.", ".self_attn.") converted[new_key] = value - + # === Cross-Attention === elif ".cross_attn.q_linear." in key: new_key = key.replace(".cross_attn.q_linear.", ".cross_attn.to_q.") converted[new_key] = value - + elif ".cross_attn.kv_linear.weight" in key: block_idx = key.split(".")[1] kv_weight = value kv_bias_key = key.replace(".weight", ".bias") kv_bias = source_weights.get(kv_bias_key) - + (k, v), (k_bias, v_bias) = split_kv(kv_weight, kv_bias) - + converted[f"blocks.{block_idx}.cross_attn.to_k.weight"] = k converted[f"blocks.{block_idx}.cross_attn.to_v.weight"] = v - + if k_bias is not None: converted[f"blocks.{block_idx}.cross_attn.to_k.bias"] = k_bias converted[f"blocks.{block_idx}.cross_attn.to_v.bias"] = v_bias - + processed_keys.add(key) if kv_bias is not None: processed_keys.add(kv_bias_key) - + elif ".cross_attn.kv_linear.bias" in key: continue - + elif ".cross_attn.proj." in key: new_key = key.replace(".cross_attn.proj.", ".cross_attn.to_out.") converted[new_key] = value - + elif ".cross_attn.q_norm." in key or ".cross_attn.k_norm." in key: converted[key] = value - + # === Final Layer (must come BEFORE general transformer block patterns) === elif key.startswith("final_layer.adaLN_modulation.1."): new_key = key.replace("final_layer.adaLN_modulation.1.", "final_layer.adaln_linear.") converted[new_key] = value - + # === Transformer Block AdaLN === elif ".adaLN_modulation.1." in key: new_key = key.replace(".adaLN_modulation.1.", ".adaln_linear_1.") converted[new_key] = value - + # === Transformer Block Normalization === elif ".mod_norm_attn." in key or ".mod_norm_ffn." in key: continue - + elif ".pre_crs_attn_norm.weight" in key: new_key = key.replace(".pre_crs_attn_norm.", ".norm_cross.") converted[new_key] = value - + elif ".pre_crs_attn_norm.bias" in key: new_key = key.replace(".pre_crs_attn_norm.", ".norm_cross.") converted[new_key] = value - + # === FFN (SwiGLU) === elif ".ffn.w1." in key or ".ffn.w2." in key or ".ffn.w3." in key: converted[key] = value - + elif key.startswith("final_layer.norm_final."): continue - + elif key.startswith("final_layer.linear."): new_key = key.replace("final_layer.linear.", "final_layer.proj.") converted[new_key] = value - + else: print(f" ⚠️ Unknown key: {key}") converted[key] = value - + return converted def validate_conversion(original: dict, converted: dict) -> bool: """Validate that conversion preserved all parameters correctly.""" print("\n Validating conversion...") - + orig_count = sum(p.numel() for p in original.values()) conv_count = sum(p.numel() for p in converted.values()) - + dropped_count = 0 for key, value in original.items(): if ".mod_norm_attn." in key or ".mod_norm_ffn." in key: dropped_count += value.numel() elif "final_layer.norm_final." in key: dropped_count += value.numel() - + expected_conv_count = orig_count - dropped_count - + print(f" Original parameters: {orig_count:,}") print(f" Converted parameters: {conv_count:,}") print(f" Dropped parameters (norms without params): {dropped_count:,}") - + if conv_count != expected_conv_count: print(f" ⚠️ Parameter count mismatch!") return False - + print(f" ✓ Parameter count matches") - + # Verify QKV/KV splits print("\n Verifying QKV/KV splits...") num_blocks = 48 - + for i in range(num_blocks): orig_qkv_weight = original.get(f"blocks.{i}.attn.qkv.weight") if orig_qkv_weight is not None: @@ -243,7 +243,7 @@ def validate_conversion(original: dict, converted: dict) -> bool: if not torch.allclose(orig_qkv_weight, reconstructed): print(f" ❌ QKV weight mismatch in block {i}") return False - + orig_kv_weight = original.get(f"blocks.{i}.cross_attn.kv_linear.weight") if orig_kv_weight is not None: conv_k = converted[f"blocks.{i}.cross_attn.to_k.weight"] @@ -252,7 +252,7 @@ def validate_conversion(original: dict, converted: dict) -> bool: if not torch.allclose(orig_kv_weight, reconstructed): print(f" ❌ KV weight mismatch in block {i}") return False - + print(f" ✓ All splits verified successfully") return True @@ -261,7 +261,7 @@ def copy_component(source_dir: Path, output_dir: Path, component: str, mapping: """Copy a component directory, optionally with name mapping.""" source_name = mapping.get(component, component) if mapping else component source_path = source_dir / source_name - + if source_path.exists(): output_path = output_dir / component if output_path.exists(): @@ -294,10 +294,10 @@ def update_transformer_config(transformer_dir: Path): if not config_path.exists(): print(" ⚠️ Transformer config not found, skipping") return - + with open(config_path, 'r') as f: config = json.load(f) - + if '_class_name' in config: old_class = config['_class_name'] config['_class_name'] = 'LongCatTransformer3DModel' @@ -305,15 +305,15 @@ def update_transformer_config(transformer_dir: Path): else: config['_class_name'] = 'LongCatTransformer3DModel' print(f" Added _class_name: LongCatTransformer3DModel") - + # Fix num_heads -> num_attention_heads for FastVideo compatibility if 'num_heads' in config and 'num_attention_heads' not in config: config['num_attention_heads'] = config.pop('num_heads') print(f" Updated num_heads → num_attention_heads") - + with open(config_path, 'w') as f: json.dump(config, f, indent=2) - + print(" ✓ Transformer config updated") @@ -321,13 +321,14 @@ def update_transformer_config(transformer_dir: Path): # LoRA Conversion Functions # ============================================================================ + def parse_lora_key(key: str) -> tuple[str, str]: """Parse LongCat LoRA key into module path and weight type.""" if key.startswith("lora___lorahyphen___"): key = key[len("lora___lorahyphen___"):] - + key = key.replace("___lorahyphen___", ".") - + if ".lora_down.weight" in key: return key.replace(".lora_down.weight", ""), "lora_down.weight" elif ".lora_up.weight" in key: @@ -340,7 +341,7 @@ def parse_lora_key(key: str) -> tuple[str, str]: return key.replace(".alpha_scale", ""), "alpha_scale" elif ".lora_alpha" in key: return key.replace(".lora_alpha", ""), "lora_alpha" - + raise ValueError(f"Unknown LoRA key format: {key}") @@ -350,43 +351,41 @@ def map_lora_module(module_path: str) -> list[tuple[str, str]]: match = re.match(r"blocks\.(\d+)\.attn\.qkv", module_path) if match: b = match.group(1) - return [(f"blocks.{b}.self_attn.to_q", "q"), - (f"blocks.{b}.self_attn.to_k", "k"), + return [(f"blocks.{b}.self_attn.to_q", "q"), (f"blocks.{b}.self_attn.to_k", "k"), (f"blocks.{b}.self_attn.to_v", "v")] - + # Self-attention output match = re.match(r"blocks\.(\d+)\.attn\.proj", module_path) if match: return [(f"blocks.{match.group(1)}.self_attn.to_out", "single")] - + # Cross-attention Q match = re.match(r"blocks\.(\d+)\.cross_attn\.q_linear", module_path) if match: return [(f"blocks.{match.group(1)}.cross_attn.to_q", "single")] - + # Cross-attention KV → K, V match = re.match(r"blocks\.(\d+)\.cross_attn\.kv_linear", module_path) if match: b = match.group(1) - return [(f"blocks.{b}.cross_attn.to_k", "k"), - (f"blocks.{b}.cross_attn.to_v", "v")] - + return [(f"blocks.{b}.cross_attn.to_k", "k"), (f"blocks.{b}.cross_attn.to_v", "v")] + # FFN match = re.match(r"blocks\.(\d+)\.ffn\.(w[123])", module_path) if match: return [(f"blocks.{match.group(1)}.ffn.{match.group(2)}", "single")] - + # AdaLN modulation match = re.match(r"blocks\.(\d+)\.adaLN_modulation\.1", module_path) if match: return [(f"blocks.{match.group(1)}.adaln_linear_1", "single")] - + # Final layer if module_path == "final_layer.adaLN_modulation.1": return [("final_layer.adaln_linear", "single")] if module_path == "final_layer.linear": return [("final_layer.proj", "single")] - + raise ValueError(f"Unknown LoRA module: {module_path}") @@ -394,9 +393,9 @@ def convert_lora_weights(source_weights: dict[str, torch.Tensor], lora_name: str """Convert LongCat LoRA to FastVideo format.""" print(f" Converting {lora_name}...") print(f" Source keys: {len(source_weights)}") - + converted = OrderedDict() - + # Group by module modules = {} for key in source_weights.keys(): @@ -407,24 +406,25 @@ def convert_lora_weights(source_weights: dict[str, torch.Tensor], lora_name: str modules[module_path][weight_type] = key except ValueError: continue - + # Process each module for module_path, weight_keys in modules.items(): try: targets = map_lora_module(module_path) except ValueError: continue - + # Get alpha_scale if present (defaults to 1.0 if missing) alpha_scale = 1.0 if "alpha_scale" in weight_keys: alpha_scale_tensor = source_weights[weight_keys["alpha_scale"]] - alpha_scale = alpha_scale_tensor.item() if alpha_scale_tensor.numel() == 1 else float(alpha_scale_tensor.mean()) - + alpha_scale = alpha_scale_tensor.item() if alpha_scale_tensor.numel() == 1 else float( + alpha_scale_tensor.mean()) + # Handle lora_down (lora_A) if "lora_down.weight" in weight_keys: lora_down = source_weights[weight_keys["lora_down.weight"]] - + if len(targets) == 1: converted[f"{targets[0][0]}.lora_A"] = lora_down # Compute alpha from alpha_scale and rank @@ -436,34 +436,34 @@ def convert_lora_weights(source_weights: dict[str, torch.Tensor], lora_name: str n = len(targets) rank = lora_down.shape[0] // n for i, (path, _) in enumerate(targets): - converted[f"{path}.lora_A"] = lora_down[i*rank:(i+1)*rank, :] + converted[f"{path}.lora_A"] = lora_down[i * rank:(i + 1) * rank, :] # Compute alpha from alpha_scale and rank for each split alpha = alpha_scale * rank converted[f"{path}.lora_alpha"] = torch.tensor(alpha, dtype=torch.float32) - + # Handle lora_up (lora_B) - may have multiple blocks lora_up_blocks = [] i = 0 while f"lora_up.blocks.{i}.weight" in weight_keys: lora_up_blocks.append(source_weights[weight_keys[f"lora_up.blocks.{i}.weight"]]) i += 1 - + if lora_up_blocks: # Multi-block LoRA: construct block-diagonal lora_B # This is equivalent to the multi-block computation without modifying fastvideo n_blocks = len(lora_up_blocks) out_per_block, rank_per_block = lora_up_blocks[0].shape # e.g., [4096, 128] - + if len(targets) == 1: # Single layer with multi-block: create block-diagonal matrix total_out = out_per_block * n_blocks total_rank = rank_per_block * n_blocks lora_B_blockdiag = torch.zeros(total_out, total_rank, dtype=lora_up_blocks[0].dtype) - + for i in range(n_blocks): - lora_B_blockdiag[i*out_per_block:(i+1)*out_per_block, - i*rank_per_block:(i+1)*rank_per_block] = lora_up_blocks[i] - + lora_B_blockdiag[i * out_per_block:(i + 1) * out_per_block, + i * rank_per_block:(i + 1) * rank_per_block] = lora_up_blocks[i] + converted[f"{targets[0][0]}.lora_B"] = lora_B_blockdiag # Note: rank for alpha calculation should be total_rank (will be computed from lora_A.shape[0]) else: @@ -481,10 +481,10 @@ def convert_lora_weights(source_weights: dict[str, torch.Tensor], lora_name: str n = len(targets) out_dim = lora_up.shape[0] // n for i, (path, _) in enumerate(targets): - converted[f"{path}.lora_B"] = lora_up[i*out_dim:(i+1)*out_dim, :] + converted[f"{path}.lora_B"] = lora_up[i * out_dim:(i + 1) * out_dim, :] else: continue - + print(f" Output keys: {len(converted)} (including lora_alpha)") # Count how many lora_alpha values were added alpha_count = sum(1 for k in converted.keys() if "lora_alpha" in k) @@ -498,21 +498,21 @@ def convert_loras(source_dir: Path, output_dir: Path) -> bool: if not lora_source.exists(): print(" No LoRA directory found, skipping") return False - + lora_files = list(lora_source.glob("*.safetensors")) if not lora_files: print(" No LoRA files found, skipping") return False - + print(f" Found {len(lora_files)} LoRA file(s)") - + # Map LoRA filenames to subdirectory names for FastVideo compatibility # Each LoRA gets its own directory under lora/ lora_subdir_mapping = { "cfg_step_lora.safetensors": "distilled", "refinement_lora.safetensors": "refinement", } - + for lora_file in lora_files: try: # Determine output subdirectory - use mapping if available, otherwise generic name @@ -521,34 +521,32 @@ def convert_loras(source_dir: Path, output_dir: Path) -> bool: else: # For unknown LoRAs, create subdirectory based on filename lora_subdir_name = lora_file.stem - + lora_output = output_dir / "lora" / lora_subdir_name lora_output.mkdir(parents=True, exist_ok=True) - + # Load source_weights = load_file(str(lora_file)) - + # Convert converted = convert_lora_weights(source_weights, lora_file.stem) - + # Save output_file = lora_output / lora_file.name save_file(converted, str(output_file)) - + size_mb = output_file.stat().st_size / (1024**2) print(f" ✓ {lora_file.name} → lora/{lora_subdir_name}/ ({size_mb:.1f} MB)") - + except Exception as e: print(f" ❌ Failed to convert {lora_file.name}: {e}") return False - + return True def main(): - parser = argparse.ArgumentParser( - description="Convert LongCat weights to FastVideo native format" - ) + parser = argparse.ArgumentParser(description="Convert LongCat weights to FastVideo native format") parser.add_argument( "--source", type=str, @@ -566,58 +564,58 @@ def main(): action="store_true", help="Run validation after conversion", ) - + args = parser.parse_args() - + source_dir = Path(args.source) output_dir = Path(args.output) - + # Check source directory if not source_dir.exists(): print(f"❌ Error: Source directory not found: {source_dir}") return 1 - + # Check for dit/transformer directory (original uses 'dit', we output to 'transformer') transformer_source = source_dir / "dit" if not transformer_source.exists(): print(f"❌ Error: DiT directory not found in source") return 1 - + print("=" * 60) print("LongCat → FastVideo Native Conversion") print("=" * 60) print(f"Source: {source_dir}") print(f"Output: {output_dir}") print() - + # Step 1: Convert transformer weights print("[Step 1/4] Converting transformer weights...") - + # Load source weights shard_files = sorted(glob.glob(str(transformer_source / "*.safetensors"))) if not shard_files: print(f"❌ Error: No safetensors files found in {transformer_source}") return 1 - + print(f" Found {len(shard_files)} shard(s)") source_weights = {} for shard_file in shard_files: print(f" Loading {Path(shard_file).name}...") source_weights.update(load_file(shard_file)) - + print(f" Loaded {len(source_weights)} parameters") - + # Convert converted_weights = convert_transformer_weights(source_weights) print(f" Converted to {len(converted_weights)} parameters") - + # Validate if requested if args.validate: if not validate_conversion(source_weights, converted_weights): print("\n❌ Validation failed!") return 1 print("\n✓ Validation passed!") - + # Save transformer_output = output_dir / "transformer" transformer_output.mkdir(parents=True, exist_ok=True) @@ -627,35 +625,35 @@ def main(): size_gb = output_file.stat().st_size / (1024**3) print(f" ✓ Saved ({size_gb:.2f} GB)") print() - + # Step 2: Copy other components print("[Step 2/5] Copying other components...") output_dir.mkdir(parents=True, exist_ok=True) - + components = ["vae", "text_encoder", "tokenizer", "scheduler"] for component in components: copy_component(source_dir, output_dir, component) - + print() - + # Step 3: Convert LoRA weights print("[Step 3/5] Converting LoRA weights...") convert_loras(source_dir, output_dir) print() - + # Step 4: Update transformer config print("[Step 4/5] Updating transformer config...") - + # Copy config.json from source if exists source_config = transformer_source / "config.json" output_config = transformer_output / "config.json" if source_config.exists(): shutil.copy(source_config, output_config) print(f" Copied config.json") - + update_transformer_config(transformer_output) print() - + # Step 5: Create model_index.json print("[Step 5/5] Creating model_index.json...") model_index_path = output_dir / "model_index.json" @@ -663,7 +661,7 @@ def main(): json.dump(create_model_index(), f, indent=2) print(f" ✓ Created {model_index_path}") print() - + print("=" * 60) print("✓ Conversion Complete!") print("=" * 60) @@ -700,7 +698,7 @@ def main(): print(" guidance_scale=1.0)") print() print() - + return 0 diff --git a/scripts/checkpoint_conversion/pt_to_safetensors.py b/scripts/checkpoint_conversion/pt_to_safetensors.py index 25f0909c9..7f7ad2ead 100644 --- a/scripts/checkpoint_conversion/pt_to_safetensors.py +++ b/scripts/checkpoint_conversion/pt_to_safetensors.py @@ -22,9 +22,7 @@ def convert_pt_to_safetensors( raise FileNotFoundError(f"Input file not found: {input_path}") if output_path.exists() and not force: - raise FileExistsError( - f"Output file already exists: {output_path}. Use --force to overwrite." - ) + raise FileExistsError(f"Output file already exists: {output_path}. Use --force to overwrite.") checkpoint = torch.load(input_path, map_location="cpu") @@ -48,59 +46,30 @@ def convert_pt_to_safetensors( raise TypeError(f"Expected a dict state_dict, got {type(state_dict)}") if skip_patterns: - state_dict = { - k: v - for k, v in state_dict.items() - if not any(pat in k for pat in skip_patterns) - } + state_dict = {k: v for k, v in state_dict.items() if not any(pat in k for pat in skip_patterns)} output_path.parent.mkdir(parents=True, exist_ok=True) save_file(state_dict, str(output_path)) def main(): - parser = argparse.ArgumentParser( - description="Convert a PyTorch checkpoint (.pt) to safetensors." - ) - parser.add_argument( - "input", - type=str, - help="Path to input .pt checkpoint file" - ) - parser.add_argument( - "output", - type=str, - help="Path to output .safetensors file" - ) - parser.add_argument( - "--key", - type=str, - default=None, - help="Optional key to extract from checkpoint dict (e.g. 'model_state_dict')" - ) - parser.add_argument( - "--force", - action="store_true", - help="Overwrite output file if it exists" - ) - parser.add_argument( - "--skip-pattern", - action="append", - dest="skip_patterns", - help="Parameter name patterns to skip (can be used multiple times)" - ) + parser = argparse.ArgumentParser(description="Convert a PyTorch checkpoint (.pt) to safetensors.") + parser.add_argument("input", type=str, help="Path to input .pt checkpoint file") + parser.add_argument("output", type=str, help="Path to output .safetensors file") + parser.add_argument("--key", + type=str, + default=None, + help="Optional key to extract from checkpoint dict (e.g. 'model_state_dict')") + parser.add_argument("--force", action="store_true", help="Overwrite output file if it exists") + parser.add_argument("--skip-pattern", + action="append", + dest="skip_patterns", + help="Parameter name patterns to skip (can be used multiple times)") args = parser.parse_args() - convert_pt_to_safetensors( - args.input, - args.output, - args.key, - args.force, - args.skip_patterns - ) + convert_pt_to_safetensors(args.input, args.output, args.key, args.force, args.skip_patterns) if __name__ == "__main__": main() - diff --git a/scripts/checkpoint_conversion/push_magi_human_to_hf.py b/scripts/checkpoint_conversion/push_magi_human_to_hf.py index f35211f12..ba212e941 100644 --- a/scripts/checkpoint_conversion/push_magi_human_to_hf.py +++ b/scripts/checkpoint_conversion/push_magi_human_to_hf.py @@ -35,10 +35,8 @@ def _validate_local_dir(local_dir: Path) -> None: required = ["model_index.json", "transformer"] missing = [r for r in required if not (local_dir / r).exists()] if missing: - sys.exit( - f"Error: {local_dir} is missing {missing}. Run " - f"convert_magi_human_to_diffusers.py first." - ) + sys.exit(f"Error: {local_dir} is missing {missing}. Run " + f"convert_magi_human_to_diffusers.py first.") def main() -> None: @@ -73,15 +71,9 @@ def main() -> None: local_dir = Path(args.local_dir).resolve() _validate_local_dir(local_dir) - token = ( - os.environ.get("HF_TOKEN") - or os.environ.get("HUGGINGFACE_HUB_TOKEN") - or os.environ.get("HF_API_KEY") - ) + token = (os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get("HF_API_KEY")) if not token: - sys.exit( - "Error: no HF token in env (set HF_TOKEN / HUGGINGFACE_HUB_TOKEN / HF_API_KEY)." - ) + sys.exit("Error: no HF token in env (set HF_TOKEN / HUGGINGFACE_HUB_TOKEN / HF_API_KEY).") api = HfApi() me = api.whoami(token=token) diff --git a/scripts/checkpoint_conversion/stable_audio_to_diffusers.py b/scripts/checkpoint_conversion/stable_audio_to_diffusers.py index eee392299..4435d6170 100644 --- a/scripts/checkpoint_conversion/stable_audio_to_diffusers.py +++ b/scripts/checkpoint_conversion/stable_audio_to_diffusers.py @@ -104,9 +104,9 @@ def _rename_layernorm_keys(state: dict[str, Any]) -> dict[str, Any]: out: dict[str, Any] = {} for k, v in state.items(): if k.endswith(".gamma"): - out[k[: -len(".gamma")] + ".weight"] = v + out[k[:-len(".gamma")] + ".weight"] = v elif k.endswith(".beta"): - out[k[: -len(".beta")] + ".bias"] = v + out[k[:-len(".beta")] + ".bias"] = v else: out[k] = v return out @@ -134,9 +134,7 @@ def _split_state_dict(src_safetensors: Path) -> dict[str, dict[str, Any]]: return buckets -def _detect_projection_flags( - diff_cfg: dict[str, Any], - dit_state: dict[str, Any]) -> tuple[bool, bool]: +def _detect_projection_flags(diff_cfg: dict[str, Any], dit_state: dict[str, Any]) -> tuple[bool, bool]: """Decide `project_cond_tokens` / `project_global_cond` by comparing the actual `to_cond_embed.0` / `to_global_embed.0` weight shapes in the DiT state dict against `cond_token_dim` / `global_cond_dim`. @@ -243,9 +241,8 @@ def convert(src: str, dst: str) -> None: raise FileNotFoundError(f"Expected monolithic safetensors at {monolithic}") cfg_path = src_dir / "model_config.json" if not cfg_path.is_file(): - raise FileNotFoundError( - f"Expected model_config.json at {cfg_path} (Stable Audio's authoritative " - f"per-component arch config).") + raise FileNotFoundError(f"Expected model_config.json at {cfg_path} (Stable Audio's authoritative " + f"per-component arch config).") with open(cfg_path) as f: model_config = json.load(f) @@ -257,8 +254,7 @@ def convert(src: str, dst: str) -> None: print("\n[3/4] Writing per-component configs + safetensors:") component_cfgs = _component_config(model_config) - project_cond, project_glob = _detect_projection_flags(component_cfgs["transformer"], - buckets["transformer"]) + project_cond, project_glob = _detect_projection_flags(component_cfgs["transformer"], buckets["transformer"]) component_cfgs["transformer"]["project_cond_tokens"] = project_cond component_cfgs["transformer"]["project_global_cond"] = project_glob for name in ("transformer", "vae", "conditioner"): @@ -284,7 +280,8 @@ def convert(src: str, dst: str) -> None: def main() -> None: parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) - parser.add_argument("--src", required=True, + parser.add_argument("--src", + required=True, help="HF repo id or local directory containing model.safetensors + model_config.json") parser.add_argument("--dst", required=True, help="Output directory for the Diffusers-format repo") args = parser.parse_args() diff --git a/scripts/checkpoint_conversion/validate_longcat_weights.py b/scripts/checkpoint_conversion/validate_longcat_weights.py index d134ac47d..d0416dc21 100644 --- a/scripts/checkpoint_conversion/validate_longcat_weights.py +++ b/scripts/checkpoint_conversion/validate_longcat_weights.py @@ -10,7 +10,7 @@ def validate_components(model_path): print("=" * 60) print("VALIDATING COMPONENTS") print("=" * 60) - + components = { "tokenizer": ["special_tokens_map.json", "tokenizer_config.json"], "text_encoder": ["config.json", "model.safetensors.index.json"], @@ -18,17 +18,17 @@ def validate_components(model_path): "scheduler": ["scheduler_config.json"], "transformer": ["config.json", "diffusion_pytorch_model.safetensors.index.json"] } - + all_valid = True for component, required_files in components.items(): component_path = os.path.join(model_path, component) print(f"\n{component}:") - + if not os.path.exists(component_path): print(f" ✗ Directory not found") all_valid = False continue - + for req_file in required_files: file_path = os.path.join(component_path, req_file) exists = os.path.exists(file_path) @@ -36,7 +36,7 @@ def validate_components(model_path): print(f" {symbol} {req_file}") if not exists: all_valid = False - + return all_valid @@ -45,70 +45,70 @@ def validate_dit_weights(dit_path): print("\n" + "=" * 60) print("VALIDATING DiT WEIGHTS") print("=" * 60) - + # Load config config_path = os.path.join(dit_path, "config.json") with open(config_path) as f: config = json.load(f) - + hidden_size = config["hidden_size"] depth = config["depth"] num_heads = config["num_heads"] - + print(f"\nArchitecture:") print(f" - hidden_size: {hidden_size}") print(f" - depth: {depth}") print(f" - num_heads: {num_heads}") - + # Load weight index index_path = os.path.join(dit_path, "diffusion_pytorch_model.safetensors.index.json") with open(index_path) as f: index = json.load(f) - + weight_map = index["weight_map"] all_keys = list(weight_map.keys()) - + print(f"\nWeight statistics:") print(f" - Total keys: {len(all_keys)}") print(f" - Total size: {index['metadata']['total_size'] / 1e9:.2f} GB") - + # Check structure embedder_keys = [k for k in all_keys if 'embedder' in k] block_keys = [k for k in all_keys if k.startswith('blocks.')] final_keys = [k for k in all_keys if k.startswith('final_layer.')] - + print(f"\nKey distribution:") print(f" - Embedder layers: {len(embedder_keys)}") print(f" - Transformer blocks: {len(block_keys)}") print(f" - Final layer: {len(final_keys)}") - + # Verify all blocks present block_nums = set() for key in block_keys: if key.startswith('blocks.'): block_num = int(key.split('.')[1]) block_nums.add(block_num) - + expected_blocks = set(range(depth)) missing_blocks = expected_blocks - block_nums - + if missing_blocks: print(f"\n✗ Missing blocks: {sorted(missing_blocks)}") return False else: print(f"\n✓ All {depth} blocks present (0-{depth-1})") - + # Sample weights first_shard = os.path.join(dit_path, "diffusion_pytorch_model-00001-of-00006.safetensors") print(f"\nSampling weights from first shard:") - + with safe_open(first_shard, framework="pt", device="cpu") as f: sample_keys = [k for k in f.keys() if k in all_keys][:5] for key in sample_keys: tensor = f.get_tensor(key) print(f" - {key}") print(f" Shape: {tuple(tensor.shape)}, Dtype: {tensor.dtype}") - + return True @@ -117,34 +117,34 @@ def validate_shapes(dit_path): print("\n" + "=" * 60) print("VALIDATING WEIGHT SHAPES") print("=" * 60) - + # Load config config_path = os.path.join(dit_path, "config.json") with open(config_path) as f: config = json.load(f) - + hidden_size = config["hidden_size"] num_heads = config["num_heads"] head_dim = hidden_size // num_heads adaln_dim = config.get("adaln_tembed_dim", 512) mlp_ratio = config.get("mlp_ratio", 4) - + # Calculate FFN hidden_dim using SwiGLU formula from blocks.py # hidden_dim = int(2 * (hidden_size * mlp_ratio) / 3) # rounded to multiple_of=256 multiple_of = 256 ffn_hidden = int(2 * hidden_size * mlp_ratio / 3) ffn_hidden = multiple_of * ((ffn_hidden + multiple_of - 1) // multiple_of) - + # Expected shapes expected = { "x_embedder.proj.weight": (hidden_size, 16, 1, 2, 2), - "x_embedder.proj.bias": (hidden_size,), + "x_embedder.proj.bias": (hidden_size, ), "t_embedder.mlp.0.weight": (adaln_dim, 256), "t_embedder.mlp.2.weight": (adaln_dim, adaln_dim), "y_embedder.y_proj.0.weight": (hidden_size, 4096), "blocks.0.attn.qkv.weight": (3 * hidden_size, hidden_size), - "blocks.0.attn.q_norm.weight": (head_dim,), + "blocks.0.attn.q_norm.weight": (head_dim, ), "blocks.0.attn.proj.weight": (hidden_size, hidden_size), "blocks.0.cross_attn.q_linear.weight": (hidden_size, hidden_size), "blocks.0.cross_attn.kv_linear.weight": (2 * hidden_size, hidden_size), @@ -152,23 +152,23 @@ def validate_shapes(dit_path): "blocks.0.adaLN_modulation.1.weight": (6 * hidden_size, adaln_dim), "final_layer.linear.weight": (64, hidden_size), } - + # Load and check first_shard = os.path.join(dit_path, "diffusion_pytorch_model-00001-of-00006.safetensors") - + all_valid = True with safe_open(first_shard, framework="pt", device="cpu") as f: for key, expected_shape in expected.items(): if key in f.keys(): tensor = f.get_tensor(key) actual_shape = tuple(tensor.shape) - + if actual_shape == expected_shape: print(f"✓ {key}: {actual_shape}") else: print(f"✗ {key}: expected {expected_shape}, got {actual_shape}") all_valid = False - + return all_valid @@ -177,19 +177,18 @@ def validate_model_index(model_path): print("\n" + "=" * 60) print("VALIDATING MODEL INDEX") print("=" * 60) - + model_index_path = os.path.join(model_path, "model_index.json") - + if not os.path.exists(model_index_path): print("✗ model_index.json not found") return False - + with open(model_index_path) as f: index = json.load(f) - - required_keys = ["_class_name", "workload_type", "tokenizer", "text_encoder", - "vae", "scheduler", "transformer"] - + + required_keys = ["_class_name", "workload_type", "tokenizer", "text_encoder", "vae", "scheduler", "transformer"] + all_valid = True for key in required_keys: if key in index: @@ -197,45 +196,42 @@ def validate_model_index(model_path): else: print(f"✗ {key}: missing") all_valid = False - + return all_valid if __name__ == "__main__": parser = argparse.ArgumentParser(description="Validate LongCat weights for FastVideo") - parser.add_argument("--model-path", type=str, required=True, - help="Path to LongCat model directory") - parser.add_argument("--check-shapes", action="store_true", - help="Also validate weight shapes (slower)") - + parser.add_argument("--model-path", type=str, required=True, help="Path to LongCat model directory") + parser.add_argument("--check-shapes", action="store_true", help="Also validate weight shapes (slower)") + args = parser.parse_args() - + print(f"\nValidating: {args.model_path}\n") - + # Run validations components_valid = validate_components(args.model_path) model_index_valid = validate_model_index(args.model_path) dit_valid = validate_dit_weights(os.path.join(args.model_path, "transformer")) - + if args.check_shapes: shapes_valid = validate_shapes(os.path.join(args.model_path, "transformer")) else: shapes_valid = True print("\nSkipping shape validation (use --check-shapes to enable)") - + # Summary print("\n" + "=" * 60) print("VALIDATION SUMMARY") print("=" * 60) - + all_valid = components_valid and model_index_valid and dit_valid and shapes_valid - + if all_valid: print("✓ All validations passed!") print("✓ Model ready for FastVideo") else: print("✗ Some validations failed") print("✗ Please check errors above") - - print("=" * 60) + print("=" * 60) diff --git a/scripts/checkpoint_conversion/wan_to_diffusers.py b/scripts/checkpoint_conversion/wan_to_diffusers.py index ea52ed687..2f6e1a626 100644 --- a/scripts/checkpoint_conversion/wan_to_diffusers.py +++ b/scripts/checkpoint_conversion/wan_to_diffusers.py @@ -6,133 +6,72 @@ import re from collections import OrderedDict _param_names_mapping: dict = { - r"^text_embedding\.0\.(.*)$": - r"condition_embedder.text_embedder.linear_1.\1", - r"^text_embedding\.2\.(.*)$": - r"condition_embedder.text_embedder.linear_2.\1", - r"^time_embedding\.0\.(.*)$": - r"condition_embedder.time_embedder.linear_1.\1", - r"^time_embedding\.2\.(.*)$": - r"condition_embedder.time_embedder.linear_2.\1", - r"^time_projection\.1\.(.*)$": - r"condition_embedder.time_proj.\1", - r"^img_emb\.proj\.0\.(.*)$": - r"condition_embedder.image_embedder.norm1.\1", - r"^img_emb\.proj\.1\.(.*)$": - r"condition_embedder.image_embedder.ff.net.0.proj.\1", - r"^img_emb\.proj\.3\.(.*)$": - r"condition_embedder.image_embedder.ff.net.2.\1", - r"^img_emb\.proj\.4\.(.*)$": - r"condition_embedder.image_embedder.norm2.\1", - r"^head\.modulation": - r"scale_shift_table", - r"^head\.head\.(.*)$": - r"proj_out.\1", - r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": - r"blocks.\1.attn1.to_q.\2", - r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": - r"blocks.\1.attn1.to_k.\2", - r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": - r"blocks.\1.attn1.to_v.\2", - r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": - r"blocks.\1.attn1.to_out.0.\2", - r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": - r"blocks.\1.attn1.norm_q.\2", - r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": - r"blocks.\1.attn1.norm_k.\2", - r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": - r"blocks.\1.attn2.to_q.\2", - r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": - r"blocks.\1.attn2.to_k.\2", - r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$": - r"blocks.\1.attn2.add_k_proj.\2", - r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": - r"blocks.\1.attn2.to_v.\2", - r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$": - r"blocks.\1.attn2.add_v_proj.\2", - r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": - r"blocks.\1.attn2.to_out.0.\2", - r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": - r"blocks.\1.attn2.norm_q.\2", - r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": - r"blocks.\1.attn2.norm_k.\2", - r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$": - r"blocks.\1.attn2.norm_added_k.\2", - r"^blocks\.(\d+)\.ffn\.0\.(.*)$": - r"blocks.\1.ffn.net.0.proj.\2", - r"^blocks\.(\d+)\.ffn\.2\.(.*)$": - r"blocks.\1.ffn.net.2.\2", - r"^blocks\.(\d+)\.modulation": - r"blocks.\1.scale_shift_table", - r"^blocks\.(\d+)\.norm3\.(.*)$": - r"blocks.\1.norm2.\2", + r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.linear_1.\1", + r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.linear_2.\1", + r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.linear_1.\1", + r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.linear_2.\1", + r"^time_projection\.1\.(.*)$": r"condition_embedder.time_proj.\1", + r"^img_emb\.proj\.0\.(.*)$": r"condition_embedder.image_embedder.norm1.\1", + r"^img_emb\.proj\.1\.(.*)$": r"condition_embedder.image_embedder.ff.net.0.proj.\1", + r"^img_emb\.proj\.3\.(.*)$": r"condition_embedder.image_embedder.ff.net.2.\1", + r"^img_emb\.proj\.4\.(.*)$": r"condition_embedder.image_embedder.norm2.\1", + r"^head\.modulation": r"scale_shift_table", + r"^head\.head\.(.*)$": r"proj_out.\1", + r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2", + r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2", + r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2", + r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.attn1.to_out.0.\2", + r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.attn1.norm_q.\2", + r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.attn1.norm_k.\2", + r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", + r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", + r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$": r"blocks.\1.attn2.add_k_proj.\2", + r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", + r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$": r"blocks.\1.attn2.add_v_proj.\2", + r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.0.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$": r"blocks.\1.attn2.norm_added_k.\2", + r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.net.0.proj.\2", + r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.net.2.\2", + r"^blocks\.(\d+)\.modulation": r"blocks.\1.scale_shift_table", + r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.norm2.\2", } # The following mapping has an extra 'patch_embedding' field and also contains # the 'model' prefixes _self_forcing_to_diffusers_param_names_mapping: dict = { - r"^model.patch_embedding\.(.*)$": - r"patch_embedding.\1", - r"^model.text_embedding\.0\.(.*)$": - r"condition_embedder.text_embedder.linear_1.\1", - r"^model.text_embedding\.2\.(.*)$": - r"condition_embedder.text_embedder.linear_2.\1", - r"^model.time_embedding\.0\.(.*)$": - r"condition_embedder.time_embedder.linear_1.\1", - r"^model.time_embedding\.2\.(.*)$": - r"condition_embedder.time_embedder.linear_2.\1", - r"^model.time_projection\.1\.(.*)$": - r"condition_embedder.time_proj.\1", - r"^model.img_emb\.proj\.0\.(.*)$": - r"condition_embedder.image_embedder.norm1.\1", - r"^model.img_emb\.proj\.1\.(.*)$": - r"condition_embedder.image_embedder.ff.net.0.proj.\1", - r"^model.img_emb\.proj\.3\.(.*)$": - r"condition_embedder.image_embedder.ff.net.2.\1", - r"^model.img_emb\.proj\.4\.(.*)$": - r"condition_embedder.image_embedder.norm2.\1", - r"^model.head\.modulation": - r"scale_shift_table", - r"^model.head\.head\.(.*)$": - r"proj_out.\1", - r"^model.blocks\.(\d+)\.self_attn\.q\.(.*)$": - r"blocks.\1.attn1.to_q.\2", - r"^model.blocks\.(\d+)\.self_attn\.k\.(.*)$": - r"blocks.\1.attn1.to_k.\2", - r"^model.blocks\.(\d+)\.self_attn\.v\.(.*)$": - r"blocks.\1.attn1.to_v.\2", - r"^model.blocks\.(\d+)\.self_attn\.o\.(.*)$": - r"blocks.\1.attn1.to_out.0.\2", - r"^model.blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": - r"blocks.\1.attn1.norm_q.\2", - r"^model.blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": - r"blocks.\1.attn1.norm_k.\2", - r"^model.blocks\.(\d+)\.cross_attn\.q\.(.*)$": - r"blocks.\1.attn2.to_q.\2", - r"^model.blocks\.(\d+)\.cross_attn\.k\.(.*)$": - r"blocks.\1.attn2.to_k.\2", - r"^model.blocks\.(\d+)\.cross_attn\.k_img\.(.*)$": - r"blocks.\1.attn2.add_k_proj.\2", - r"^model.blocks\.(\d+)\.cross_attn\.v\.(.*)$": - r"blocks.\1.attn2.to_v.\2", - r"^model.blocks\.(\d+)\.cross_attn\.v_img\.(.*)$": - r"blocks.\1.attn2.add_v_proj.\2", - r"^model.blocks\.(\d+)\.cross_attn\.o\.(.*)$": - r"blocks.\1.attn2.to_out.0.\2", - r"^model.blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": - r"blocks.\1.attn2.norm_q.\2", - r"^model.blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": - r"blocks.\1.attn2.norm_k.\2", - r"^model.blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$": - r"blocks.\1.attn2.norm_added_k.\2", - r"^model.blocks\.(\d+)\.ffn\.0\.(.*)$": - r"blocks.\1.ffn.net.0.proj.\2", - r"^model.blocks\.(\d+)\.ffn\.2\.(.*)$": - r"blocks.\1.ffn.net.2.\2", - r"^model.blocks\.(\d+)\.modulation": - r"blocks.\1.scale_shift_table", - r"^model.blocks\.(\d+)\.norm3\.(.*)$": - r"blocks.\1.norm2.\2", + r"^model.patch_embedding\.(.*)$": r"patch_embedding.\1", + r"^model.text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.linear_1.\1", + r"^model.text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.linear_2.\1", + r"^model.time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.linear_1.\1", + r"^model.time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.linear_2.\1", + r"^model.time_projection\.1\.(.*)$": r"condition_embedder.time_proj.\1", + r"^model.img_emb\.proj\.0\.(.*)$": r"condition_embedder.image_embedder.norm1.\1", + r"^model.img_emb\.proj\.1\.(.*)$": r"condition_embedder.image_embedder.ff.net.0.proj.\1", + r"^model.img_emb\.proj\.3\.(.*)$": r"condition_embedder.image_embedder.ff.net.2.\1", + r"^model.img_emb\.proj\.4\.(.*)$": r"condition_embedder.image_embedder.norm2.\1", + r"^model.head\.modulation": r"scale_shift_table", + r"^model.head\.head\.(.*)$": r"proj_out.\1", + r"^model.blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2", + r"^model.blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2", + r"^model.blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2", + r"^model.blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.attn1.to_out.0.\2", + r"^model.blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.attn1.norm_q.\2", + r"^model.blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.attn1.norm_k.\2", + r"^model.blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", + r"^model.blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", + r"^model.blocks\.(\d+)\.cross_attn\.k_img\.(.*)$": r"blocks.\1.attn2.add_k_proj.\2", + r"^model.blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", + r"^model.blocks\.(\d+)\.cross_attn\.v_img\.(.*)$": r"blocks.\1.attn2.add_v_proj.\2", + r"^model.blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.0.\2", + r"^model.blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2", + r"^model.blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2", + r"^model.blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$": r"blocks.\1.attn2.norm_added_k.\2", + r"^model.blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.net.0.proj.\2", + r"^model.blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.net.2.\2", + r"^model.blocks\.(\d+)\.modulation": r"blocks.\1.scale_shift_table", + r"^model.blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.norm2.\2", } state_dict = load_state_dict_from_file("checkpoints/self_forcing_dmd.pt") @@ -154,8 +93,4 @@ for k, v in state_dict.items(): new_state_dict[dummy_key] = dummy_value del state_dict -save_torch_state_dict( - new_state_dict, - "new2/", - max_shard_size="10GB" -) +save_torch_state_dict(new_state_dict, "new2/", max_shard_size="10GB") diff --git a/scripts/dataset_preparation/generate_validation_json.py b/scripts/dataset_preparation/generate_validation_json.py index e4cc2b179..e640860ae 100644 --- a/scripts/dataset_preparation/generate_validation_json.py +++ b/scripts/dataset_preparation/generate_validation_json.py @@ -31,14 +31,12 @@ def generate_merged_validation_json(input_dir, output_file): validation_data.append(validation_entry) # Create the final validation structure - validation_json = { - "data": validation_data - } + validation_json = {"data": validation_data} # Write the validation JSON to the output file with open(output_file, "w") as f: json.dump(validation_json, f, indent=2) - + print(f"Generated validation JSON with {len(validation_data)} entries and saved to {output_file}") @@ -60,4 +58,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/scripts/dataset_preparation/prepare_json_file.py b/scripts/dataset_preparation/prepare_json_file.py index b263b7053..cb29494f3 100644 --- a/scripts/dataset_preparation/prepare_json_file.py +++ b/scripts/dataset_preparation/prepare_json_file.py @@ -13,9 +13,7 @@ from tqdm import tqdm def get_video_info(video_path): """Get video information using torchvision.""" # Read video tensor (T, C, H, W) - video_tensor, _, info = torchvision.io.read_video(str(video_path), - output_format="TCHW", - pts_unit="sec") + video_tensor, _, info = torchvision.io.read_video(str(video_path), output_format="TCHW", pts_unit="sec") num_frames = video_tensor.shape[0] height = video_tensor.shape[2] @@ -39,9 +37,7 @@ def get_video_info(video_path): } -def prepare_dataset_json(folder_path, - output_name="videos2caption.json", - num_workers=None) -> None: +def prepare_dataset_json(folder_path, output_name="videos2caption.json", num_workers=None) -> None: """Prepare dataset information from a folder containing videos and prompt.txt.""" folder_path = Path(folder_path) @@ -62,9 +58,7 @@ def prepare_dataset_json(folder_path, video_paths = [line.strip() for line in f.readlines() if line.strip()] if len(prompts) != len(video_paths): - raise ValueError( - f"Number of prompts ({len(prompts)}) does not match number of videos ({len(video_paths)})" - ) + raise ValueError(f"Number of prompts ({len(prompts)}) does not match number of videos ({len(video_paths)})") # Prepare arguments for multiprocessing process_args = [folder_path / video_path for video_path in video_paths] @@ -113,22 +107,16 @@ def prepare_dataset_json(folder_path, def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description='Prepare video dataset information in JSON format') - parser.add_argument( - '--data_folder', - type=str, - required=True, - help='Path to the folder containing videos and prompt.txt') - parser.add_argument( - '--output', - type=str, - default='videos2caption.json', - help='Name of the output JSON file (default: videos2caption.json)') - parser.add_argument('--workers', - type=int, - default=32, - help='Number of worker processes (default: 16)') + parser = argparse.ArgumentParser(description='Prepare video dataset information in JSON format') + parser.add_argument('--data_folder', + type=str, + required=True, + help='Path to the folder containing videos and prompt.txt') + parser.add_argument('--output', + type=str, + default='videos2caption.json', + help='Name of the output JSON file (default: videos2caption.json)') + parser.add_argument('--workers', type=int, default=32, help='Number of worker processes (default: 16)') return parser.parse_args() diff --git a/scripts/demo_anyflow_14b.py b/scripts/demo_anyflow_14b.py index 53afad76b..a97f806aa 100644 --- a/scripts/demo_anyflow_14b.py +++ b/scripts/demo_anyflow_14b.py @@ -42,7 +42,6 @@ from pathlib import Path import torch from safetensors.torch import load_file - ANYFLOW_LOCAL = Path(os.environ.get("ANYFLOW_LOCAL", "./anyflow-14b")).expanduser() OUT_DIR = Path(os.environ.get("ANYFLOW_DEMO_OUT", "./demo_videos")).expanduser() OUT_DIR.mkdir(parents=True, exist_ok=True) @@ -52,15 +51,13 @@ DEVICE = torch.device("cuda") DTYPE = torch.bfloat16 NUM_FRAMES = 81 HEIGHT, WIDTH = 480, 832 -PROMPT = ( - "CG game concept digital art, a majestic elephant with a vibrant tusk and sleek fur " - "running swiftly towards a herd of its kind. The elephant has a calm yet determined " - "expression, with its ears flapping slightly as it moves at high speed. The herd consists " - "of several other elephants of various ages and sizes, all moving in unison. The landscape " - "is vast savanna with rolling hills, tall grasses, and scattered acacia trees. The sun " - "sets behind the horizon, casting a warm golden glow over the scene. Low-angle view, focus " - "on the elephant as it accelerates towards the herd." -) +PROMPT = ("CG game concept digital art, a majestic elephant with a vibrant tusk and sleek fur " + "running swiftly towards a herd of its kind. The elephant has a calm yet determined " + "expression, with its ears flapping slightly as it moves at high speed. The herd consists " + "of several other elephants of various ages and sizes, all moving in unison. The landscape " + "is vast savanna with rolling hills, tall grasses, and scattered acacia trees. The sun " + "sets behind the horizon, casting a warm golden glow over the scene. Low-angle view, focus " + "on the elephant as it accelerates towards the herd.") NEG_PROMPT = "blurry, low quality, distorted" # The published nvidia/AnyFlow-* checkpoints are on-policy distilled with # fuse_guidance_scale=3.0 baked into the weights — inference uses 1.0 @@ -116,14 +113,19 @@ def encode_prompts(): tok = AutoTokenizer.from_pretrained(str(ANYFLOW_LOCAL), subfolder="tokenizer", use_fast=False) enc = UMT5EncoderModel.from_pretrained( - str(ANYFLOW_LOCAL), subfolder="text_encoder", torch_dtype=DTYPE, + str(ANYFLOW_LOCAL), + subfolder="text_encoder", + torch_dtype=DTYPE, ).to(DEVICE).eval() @torch.no_grad() def encode_one(prompts): - out = tok( - prompts, padding="max_length", max_length=512, truncation=True, - return_attention_mask=True, return_tensors="pt") + out = tok(prompts, + padding="max_length", + max_length=512, + truncation=True, + return_attention_mask=True, + return_tensors="pt") ids = out.input_ids.to(DEVICE) mask = out.attention_mask.to(DEVICE) seq_lens = mask.gt(0).sum(dim=1).long() @@ -201,19 +203,13 @@ def sample(model, text_e, neg_e, nfe: int) -> torch.Tensor: t_in = t_cur.expand(B).to(DTYPE) r_in = t_next.expand(B).to(DTYPE) with set_forward_context(current_timestep=t_in, attn_metadata=None): - flow_cond = model( - hidden_states=x, encoder_hidden_states=text_e, - timestep=t_in, r_timestep=r_in) + flow_cond = model(hidden_states=x, encoder_hidden_states=text_e, timestep=t_in, r_timestep=r_in) if GUIDANCE != 1.0: - flow_uncond = model( - hidden_states=x, encoder_hidden_states=neg_e, - timestep=t_in, r_timestep=r_in) + flow_uncond = model(hidden_states=x, encoder_hidden_states=neg_e, timestep=t_in, r_timestep=r_in) flow = flow_uncond + GUIDANCE * (flow_cond - flow_uncond) else: flow = flow_cond - x = scheduler.step( - flow, sample=x, - timestep=t_cur.repeat(B), r_timestep=t_next.repeat(B)) + x = scheduler.step(flow, sample=x, timestep=t_cur.repeat(B), r_timestep=t_next.repeat(B)) print(f" NFE={nfe} sample time: {time.time() - t0:.1f}s " f"({(time.time() - t0) / nfe:.1f}s/step)") xf = x.float() @@ -229,7 +225,9 @@ def decode_to_mp4(latents: torch.Tensor, out_path: Path) -> tuple[Path, tuple]: import imageio.v3 as iio vae = AutoencoderKLWan.from_pretrained( - str(ANYFLOW_LOCAL), subfolder="vae", torch_dtype=DTYPE, + str(ANYFLOW_LOCAL), + subfolder="vae", + torch_dtype=DTYPE, ).to(DEVICE).eval() try: vae.enable_tiling() diff --git a/scripts/lora_extraction/extract_lora.py b/scripts/lora_extraction/extract_lora.py index 7b05b0559..44d8a9f43 100644 --- a/scripts/lora_extraction/extract_lora.py +++ b/scripts/lora_extraction/extract_lora.py @@ -55,18 +55,18 @@ def load_transformer_state_dict_from_safetensors(model_path: str) -> Dict[str, t """ from huggingface_hub import snapshot_download import os - + # Download or locate the model if os.path.isdir(model_path): local_path = model_path else: local_path = snapshot_download(model_path) - + # Find transformer directory transformer_dir = os.path.join(local_path, "transformer") if not os.path.isdir(transformer_dir): raise FileNotFoundError(f"Transformer directory not found at {transformer_dir}") - + # Load all safetensors files state_dict: Dict[str, torch.Tensor] = {} for fname in sorted(os.listdir(transformer_dir)): @@ -75,13 +75,14 @@ def load_transformer_state_dict_from_safetensors(model_path: str) -> Dict[str, t with safe_open(fpath, framework='pt', device='cpu') as f: for key in f.keys(): state_dict[key] = f.get_tensor(key) - + if not state_dict: raise ValueError(f"No safetensors files found in {transformer_dir}") - + LOG.info("Loaded %d keys directly from safetensors", len(state_dict)) return state_dict + # Configure minimal logging LOG = logging.getLogger("extract_lora") @@ -140,7 +141,8 @@ def load_transformer_state_dict_from_model( pipeline_attr = getattr(pipeline, "pipeline", None) transformer = getattr(pipeline_attr, "transformer", None) if pipeline_attr else None if transformer is None: - raise RuntimeError("Transformer not found in pipeline. Expected pipeline.transformer or pipeline.modules['transformer'].") + raise RuntimeError( + "Transformer not found in pipeline. Expected pipeline.transformer or pipeline.modules['transformer'].") state_dict = transformer.state_dict() @@ -257,12 +259,12 @@ def build_adapter_from_states( continue S_sqrt = torch.sqrt(S[:chosen_rank].to(torch.float32)) - U_r = U[:, :chosen_rank].to(torch.float32) # (out, r) + U_r = U[:, :chosen_rank].to(torch.float32) # (out, r) Vh_r = Vh[:chosen_rank, :].to(torch.float32) # (r, in) - lora_B = (U_r * S_sqrt.unsqueeze(0)).contiguous() # (out, r) - tmp = (Vh_r.T * S_sqrt.unsqueeze(0)).contiguous() # (in, r) - lora_A = tmp.T.contiguous() # (r, in) + lora_B = (U_r * S_sqrt.unsqueeze(0)).contiguous() # (out, r) + tmp = (Vh_r.T * S_sqrt.unsqueeze(0)).contiguous() # (in, r) + lora_A = tmp.T.contiguous() # (r, in) base_name = key[:-len(".weight")] a_key = f"{base_name}.lora_A.weight" @@ -296,8 +298,8 @@ def build_adapter_from_states( pass avg_delta = (sum(mean_deltas) / len(mean_deltas)) if mean_deltas else 0.0 - LOG.info("Extraction complete: processed_keys=%d, extracted_layers=%d, avg_abs_delta=%.6e", - len(keys), processed, avg_delta) + LOG.info("Extraction complete: processed_keys=%d, extracted_layers=%d, avg_abs_delta=%.6e", len(keys), processed, + avg_delta) return adapter_state @@ -360,7 +362,6 @@ def extract_lora_adapter( LOG.info("Loading fine-tuned model directly from safetensors: %s", ft) ft_sd = load_transformer_state_dict_from_safetensors(ft) - resume_idx = 0 adapter_existing: Dict[str, torch.Tensor] = {} if resume and checkpoint_path and checkpoint_path.exists(): diff --git a/scripts/lora_extraction/lora_inference_comparison.py b/scripts/lora_extraction/lora_inference_comparison.py index bee3c83a9..9fc757d75 100644 --- a/scripts/lora_extraction/lora_inference_comparison.py +++ b/scripts/lora_extraction/lora_inference_comparison.py @@ -64,7 +64,7 @@ def _validate_adapter(adapter: Optional[str]) -> Optional[str]: p = Path(adapter).expanduser() if not p.exists(): raise FileNotFoundError(f"Adapter not found: {p}") - + # Accept both files and directories (FastVideo expects directories for HF-style adapters) if p.is_file(): if p.suffix != ".safetensors": @@ -78,7 +78,7 @@ def _validate_adapter(adapter: Optional[str]) -> Optional[str]: raise ValueError(f"Adapter directory contains no .safetensors files: {p}") else: raise ValueError(f"Adapter must be a file or directory: {p}") - + return str(p.resolve()) @@ -151,9 +151,10 @@ def generate_with_model( raise FileNotFoundError(f"Video not found at expected path: {expected}") -def compute_metrics(output_dir: str, ft_video: str, other_video: str, num_inference_steps: int, prompt: str, compute_ssim: bool, compute_lpips: bool) -> dict: +def compute_metrics(output_dir: str, ft_video: str, other_video: str, num_inference_steps: int, prompt: str, + compute_ssim: bool, compute_lpips: bool) -> dict: results = {} - + if compute_ssim: try: from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results # type: ignore @@ -162,40 +163,40 @@ def compute_metrics(output_dir: str, ft_video: str, other_video: str, num_infere write_ssim_results(output_dir, ssim_values, ft_video, other_video, num_inference_steps, prompt) except Exception as e: logger.warning(f"SSIM computation failed: {e}") - + if compute_lpips: try: import torch import lpips import torchvision.io as tv_io - + loss_fn = lpips.LPIPS(net='alex') - + # Load videos vid1, _, _ = tv_io.read_video(ft_video, pts_unit='sec') vid2, _, _ = tv_io.read_video(other_video, pts_unit='sec') - + # Normalize to [-1, 1] vid1 = (vid1.float() / 127.5 - 1.0).permute(0, 3, 1, 2) # (T, C, H, W) vid2 = (vid2.float() / 127.5 - 1.0).permute(0, 3, 1, 2) - + lpips_scores = [] with torch.no_grad(): for frame1, frame2 in zip(vid1, vid2): score = loss_fn(frame1.unsqueeze(0), frame2.unsqueeze(0)) lpips_scores.append(float(score.item())) - + results["mean_lpips"] = sum(lpips_scores) / len(lpips_scores) - + # Write LPIPS results import json lpips_file = Path(output_dir) / f"steps{num_inference_steps}_{prompt.replace(' ', '_')[:30]}_lpips.json" with open(lpips_file, 'w') as f: json.dump({"mean_lpips": results["mean_lpips"], "lpips_per_frame": lpips_scores}, f, indent=2) - + except Exception as e: logger.warning(f"LPIPS computation failed: {e}") - + return results @@ -321,7 +322,8 @@ def compare_inference( # 3) compute metrics results = {} if compute_ssim or compute_lpips: - results = compute_metrics(output_dir, ft_video, target_video, num_inference_steps, prompt, compute_ssim, compute_lpips) + results = compute_metrics(output_dir, ft_video, target_video, num_inference_steps, prompt, compute_ssim, + compute_lpips) if results.get("mean_ssim") is not None: logger.info("Mean SSIM: %.4f", results["mean_ssim"]) if results.get("mean_lpips") is not None: diff --git a/scripts/lora_extraction/merge_lora.py b/scripts/lora_extraction/merge_lora.py index 6cd666168..f06a45618 100644 --- a/scripts/lora_extraction/merge_lora.py +++ b/scripts/lora_extraction/merge_lora.py @@ -47,10 +47,10 @@ def configure_logging(level: str = "INFO"): def fix_adapter_naming(adapter: dict) -> dict: fixed = {} renamed = 0 - + for key, tensor in adapter.items(): new_key = key - + if ".lora_A.weight" in key: base = key.replace(".lora_A.weight", "") if not base.endswith(".weight"): @@ -71,12 +71,12 @@ def fix_adapter_naming(adapter: dict) -> dict: if not base.endswith(".weight"): new_key = base + ".weight.lora_alpha" renamed += 1 - + fixed[new_key] = tensor - + if renamed > 0: LOG.info(f"Fixed {renamed} adapter key names") - + return fixed @@ -86,18 +86,18 @@ def load_adapter(adapter_path: str) -> dict: raise FileNotFoundError(f"Adapter file not found: {abs_path}") if not abs_path.endswith('.safetensors'): raise ValueError(f"Adapter must be .safetensors: {abs_path}") - + LOG.info(f"Loading adapter: {abs_path}") adapter = load_file(abs_path) file_size_mb = os.path.getsize(abs_path) / (1024 * 1024) LOG.info(f"Loaded {len(adapter)} tensors ({file_size_mb:.1f} MB)") - + return fix_adapter_naming(adapter) def group_adapter_keys(adapter: dict) -> dict: grouped = defaultdict(dict) - + for key, tensor in adapter.items(): if key.endswith(".lora_A.weight"): grouped[key.replace(".lora_A.weight", "")]["A"] = tensor @@ -107,14 +107,14 @@ def group_adapter_keys(adapter: dict) -> dict: grouped[key.replace(".lora_rank", "")]["rank"] = tensor elif key.endswith(".lora_alpha"): grouped[key.replace(".lora_alpha", "")]["alpha"] = tensor - + LOG.info(f"Grouped {len(grouped)} LoRA layers") return grouped def get_reverse_param_mapping(base_model_path: str): LOG.info("Loading base model for parameter mapping") - + pipeline_cls = get_pipeline_class_for_model(base_model_path) pipeline = pipeline_cls.from_pretrained( base_model_path, @@ -125,17 +125,17 @@ def get_reverse_param_mapping(base_model_path: str): text_encoder_cpu_offload=True, pin_cpu_memory=True, ) - + transformer = None if hasattr(pipeline, "transformer"): transformer = pipeline.transformer elif hasattr(pipeline, "modules") and isinstance(pipeline.modules, dict): if "transformer" in pipeline.modules: transformer = pipeline.modules["transformer"] - + if transformer is None: raise RuntimeError("Could not find transformer in pipeline") - + if hasattr(transformer, "reverse_param_names_mapping"): reverse_mapping = transformer.reverse_param_names_mapping elif hasattr(transformer, "config") and hasattr(transformer.config, "arch_config"): @@ -145,101 +145,96 @@ def get_reverse_param_mapping(base_model_path: str): else: param_mapping = arch_config.param_names_mapping param_names_mapping_fn = get_param_names_mapping(param_mapping) - + from diffusers import DiffusionPipeline from huggingface_hub import snapshot_download - + if not os.path.exists(base_model_path) or not os.path.isdir(base_model_path): - model_path = snapshot_download( - repo_id=base_model_path, - ignore_patterns=["*.onnx", "*.msgpack"] - ) + model_path = snapshot_download(repo_id=base_model_path, ignore_patterns=["*.onnx", "*.msgpack"]) else: model_path = base_model_path - + hf_pipeline = DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32) hf_transformer = hf_pipeline.transformer hf_sd = hf_transformer.state_dict() - + _, reverse_mapping = hf_to_custom_state_dict(hf_sd, param_names_mapping_fn) - + del hf_pipeline del hf_transformer torch.cuda.empty_cache() else: raise RuntimeError("Could not find reverse_param_names_mapping in transformer or config") - + del pipeline del transformer torch.cuda.empty_cache() - + return reverse_mapping def merge_lora_into_base(base_sd: dict, adapter: dict) -> dict: LOG.info("Merging LoRA into base weights") - + adapter_layers = group_adapter_keys(adapter) merged_sd = dict(base_sd) - + merged_count = 0 skipped_count = 0 - + for base_name, parts in adapter_layers.items(): weight_key = base_name if base_name.endswith(".weight") else base_name + ".weight" - + if weight_key not in base_sd: skipped_count += 1 continue - + if "A" not in parts or "B" not in parts: skipped_count += 1 continue - + lora_A = parts["A"].to(torch.float32) lora_B = parts["B"].to(torch.float32) base_weight = base_sd[weight_key].to(torch.float32) - + out_dim, in_dim = base_weight.shape if lora_B.shape[0] != out_dim or lora_A.shape[1] != in_dim or lora_B.shape[1] != lora_A.shape[0]: skipped_count += 1 continue - + delta = lora_B @ lora_A - + rank = int(parts.get("rank", torch.tensor([lora_A.shape[0]])).item()) if "rank" in parts else lora_A.shape[0] alpha = float(parts.get("alpha", torch.tensor([rank])).item()) if "alpha" in parts else float(rank) - + if rank != 0 and alpha != rank: delta = delta * (alpha / float(rank)) - + merged_weight = base_weight + delta merged_sd[weight_key] = merged_weight.to(base_sd[weight_key].dtype) merged_count += 1 - + LOG.info(f"Merged {merged_count} layers, skipped {skipped_count}") return merged_sd -def save_merged_model(merged_sd: dict, base_model_path: str, ft_model_path: str, output_dir: str, reverse_mapping: dict): +def save_merged_model(merged_sd: dict, base_model_path: str, ft_model_path: str, output_dir: str, + reverse_mapping: dict): output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) - + transformer_dir = output_path / "transformer" transformer_dir.mkdir(parents=True, exist_ok=True) LOG.info("Converting to HuggingFace format") hf_merged_sd = custom_to_hf_state_dict(merged_sd, reverse_mapping) LOG.info(f"Converted {len(hf_merged_sd)} parameters") - + base_path = Path(base_model_path) if not base_path.exists() or not base_path.is_dir(): from huggingface_hub import snapshot_download - base_path = Path(snapshot_download( - repo_id=base_model_path, - ignore_patterns=["*.onnx", "*.msgpack"] - )) - + base_path = Path(snapshot_download(repo_id=base_model_path, ignore_patterns=["*.onnx", "*.msgpack"])) + LOG.info("Copying model components") for component in ["scheduler", "text_encoder", "tokenizer", "vae"]: src = base_path / component @@ -248,17 +243,16 @@ def save_merged_model(merged_sd: dict, base_model_path: str, ft_model_path: str, if dst.exists(): shutil.rmtree(dst) shutil.copytree(src, dst) - + LOG.info("Copying finetuned model config") ft_path = Path(ft_model_path) if not ft_path.exists() or not ft_path.is_dir(): from huggingface_hub import snapshot_download - ft_path = Path(snapshot_download( - repo_id=ft_model_path, - allow_patterns=["model_index.json"], - ignore_patterns=["*.onnx", "*.msgpack"] - )) - + ft_path = Path( + snapshot_download(repo_id=ft_model_path, + allow_patterns=["model_index.json"], + ignore_patterns=["*.onnx", "*.msgpack"])) + ft_index = ft_path / "model_index.json" if ft_index.exists(): shutil.copy2(ft_index, output_path / "model_index.json") @@ -267,16 +261,16 @@ def save_merged_model(merged_sd: dict, base_model_path: str, ft_model_path: str, src_index = base_path / "model_index.json" if src_index.exists(): shutil.copy2(src_index, output_path / "model_index.json") - + weight_path = transformer_dir / "diffusion_pytorch_model.safetensors" LOG.info(f"Saving merged weights to {weight_path}") to_save_hf = {k: v.detach().cpu() for k, v in hf_merged_sd.items()} save_file(to_save_hf, str(weight_path)) - + config_src = base_path / "transformer" / "config.json" if config_src.exists(): shutil.copy2(config_src, transformer_dir / "config.json") - + file_size_mb = weight_path.stat().st_size / (1024 * 1024) LOG.info(f"Saved to {output_dir} ({file_size_mb:.0f} MB, {len(hf_merged_sd)} params)") @@ -298,20 +292,20 @@ def merge_lora( log_level: Logging level """ configure_logging(log_level) - + LOG.info(f"Base: {base}") LOG.info(f"Adapter: {adapter}") LOG.info(f"Output: {output}") - + reverse_mapping = get_reverse_param_mapping(base) - + LOG.info(f"Loading base model: {base}") base_sd = load_transformer_state_dict_from_model(base) LOG.info(f"Loaded { len(base_sd)} parameters") - + adapter_sd = load_adapter(adapter) merged_sd = merge_lora_into_base(base_sd, adapter_sd) - + save_merged_model(merged_sd, base, ft, output, reverse_mapping) LOG.info("Merge complete") @@ -325,7 +319,7 @@ def main(): parser.add_argument("--output", required=True, help="Output directory") parser.add_argument("--log-level", default="INFO", help="Logging level") args = parser.parse_args() - + merge_lora( base=args.base, adapter=args.adapter, diff --git a/scripts/lora_extraction/verify_lora.py b/scripts/lora_extraction/verify_lora.py index da1278306..69a0a88e3 100644 --- a/scripts/lora_extraction/verify_lora.py +++ b/scripts/lora_extraction/verify_lora.py @@ -27,64 +27,61 @@ def configure_logging(level: str = "INFO"): def load_transformer_weights(model_path: str | Path) -> dict: """Load transformer weights from model directory.""" model_path = Path(model_path) - + if not model_path.exists() or not model_path.is_dir(): from huggingface_hub import snapshot_download LOG.info(f"Downloading {model_path} from HuggingFace Hub...") - model_path = Path(snapshot_download( - repo_id=str(model_path), - ignore_patterns=["*.onnx", "*.msgpack"] - )) - + model_path = Path(snapshot_download(repo_id=str(model_path), ignore_patterns=["*.onnx", "*.msgpack"])) + transformer_dir = model_path / "transformer" - + if not transformer_dir.exists(): raise FileNotFoundError(f"Transformer directory not found: {transformer_dir}") - + weight_files = sorted(transformer_dir.glob("*.safetensors")) if not weight_files: raise FileNotFoundError(f"No safetensors files in {transformer_dir}") - + LOG.info(f"Loading {len(weight_files)} file(s) from {transformer_dir}") - + state_dict = {} for f in weight_files: if "custom" in f.name: continue state_dict.update(load_file(str(f))) - + return state_dict def compare_models(merged_sd: dict, ft_sd: dict) -> dict: """Compare merged and finetuned model weights.""" - + common_keys = set(merged_sd.keys()) & set(ft_sd.keys()) merged_only = set(merged_sd.keys()) - set(ft_sd.keys()) ft_only = set(ft_sd.keys()) - set(merged_sd.keys()) - + LOG.info(f"Common keys: {len(common_keys)}") if merged_only: LOG.warning(f"Keys only in merged: {len(merged_only)}") if ft_only: LOG.warning(f"Keys only in finetuned: {len(ft_only)}") - + results = [] for key in sorted(common_keys): merged_param = merged_sd[key] ft_param = ft_sd[key] - + if merged_param.shape != ft_param.shape: LOG.error(f"{key}: shape mismatch {merged_param.shape} vs {ft_param.shape}") continue - + diff = (merged_param.float() - ft_param.float()).abs() max_abs = diff.max().item() mean_abs = diff.mean().item() - + merged_norm = merged_param.float().norm().item() rel_mean = (mean_abs / merged_norm * 100) if merged_norm > 0 else 0 - + results.append({ "key": key, "shape": tuple(merged_param.shape), @@ -92,7 +89,7 @@ def compare_models(merged_sd: dict, ft_sd: dict) -> dict: "mean_abs": mean_abs, "rel_mean": rel_mean }) - + return results @@ -102,35 +99,37 @@ def main(): parser.add_argument("--ft", required=True, help="Finetuned model ID or path") parser.add_argument("--log-level", default="INFO", help="Logging level") args = parser.parse_args() - + configure_logging(args.log_level) - + LOG.info(f"Loading merged model: {args.merged}") merged_sd = load_transformer_weights(args.merged) LOG.info(f"Loaded {len(merged_sd)} parameters") - + LOG.info(f"Loading finetuned model: {args.ft}") ft_sd = load_transformer_weights(args.ft) LOG.info(f"Loaded {len(ft_sd)} parameters") - + LOG.info("Comparing models...") results = compare_models(merged_sd, ft_sd) - + results.sort(key=lambda x: x["max_abs"], reverse=True) - + LOG.info(f"\nTop 10 mismatches by max_abs_error:") for i, r in enumerate(results[:10], 1): LOG.info(f"{i:2d}. {r['key']}") - LOG.info(f" shape={r['shape']}, max_abs={r['max_abs']:.3e}, mean_abs={r['mean_abs']:.3e}, rel_mean={r['rel_mean']:.4f}%") - + LOG.info( + f" shape={r['shape']}, max_abs={r['max_abs']:.3e}, mean_abs={r['mean_abs']:.3e}, rel_mean={r['rel_mean']:.4f}%" + ) + overall_mean = sum(r["mean_abs"] for r in results) / len(results) overall_max = max(r["max_abs"] for r in results) - + LOG.info(f"\nOverall metrics:") LOG.info(f" Layers compared: {len(results)}") LOG.info(f" Mean(mean_abs): {overall_mean:.3e}") LOG.info(f" Max(max_abs): {overall_max:.3e}") - + if overall_mean < 1e-4: LOG.info("\nVerification PASSED: Merge is numerically accurate") else: diff --git a/scripts/ltx2_sr_alignment.py b/scripts/ltx2_sr_alignment.py index 37e12e7c3..e50e7f28f 100644 --- a/scripts/ltx2_sr_alignment.py +++ b/scripts/ltx2_sr_alignment.py @@ -56,10 +56,7 @@ def _bootstrap_fastvideo_source() -> None: # Remove editable-install finders that would otherwise win over # PYTHONPATH. uv-installed editables register a custom finder # named like ``__editable___fastvideo_0_1_7_finder``. - sys.meta_path = [ - finder for finder in sys.meta_path - if "__editable___fastvideo" not in type(finder).__module__ - ] + sys.meta_path = [finder for finder in sys.meta_path if "__editable___fastvideo" not in type(finder).__module__] # Drop pre-resolved fastvideo modules from any earlier import. for name in [k for k in sys.modules if k == "fastvideo" or k.startswith("fastvideo.")]: sys.modules.pop(name, None) @@ -71,23 +68,21 @@ def _bootstrap_fastvideo_source() -> None: _bootstrap_fastvideo_source() - import numpy as np # noqa: E402 (after bootstrap so torch picks up the right env) import torch # noqa: E402 # Pinned alignment fixture — matches basic_ltx2_upscale.py upstream. -PROMPT = ( - "A warm sunny backyard. The camera starts in a tight cinematic close-up " - "of a woman and a man in their 30s, facing each other with serious " - "expressions. The woman, emotional and dramatic, says softly, \"That's " - "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " - "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " - "then mutters defensively, \"He's just having fun.\" The camera slowly " - "pans right, revealing the grandfather in the garden wearing enormous " - "butterfly wings, waving his arms in the air like he's trying to take " - "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. " - "The woman covers her face, on the verge of tears. The tone is deadpan, " - "absurd, and quietly tragic.") +PROMPT = ("A warm sunny backyard. The camera starts in a tight cinematic close-up " + "of a woman and a man in their 30s, facing each other with serious " + "expressions. The woman, emotional and dramatic, says softly, \"That's " + "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " + "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " + "then mutters defensively, \"He's just having fun.\" The camera slowly " + "pans right, revealing the grandfather in the garden wearing enormous " + "butterfly wings, waving his arms in the air like he's trying to take " + "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. " + "The woman covers her face, on the verge of tears. The tone is deadpan, " + "absurd, and quietly tragic.") DEFAULT_MODEL = "FastVideo/LTX2-Distilled-Diffusers" DEFAULT_SEED = 10 @@ -185,8 +180,8 @@ def _run_typed(args: argparse.Namespace) -> RunResult: pipeline.preset_overrides["refine"] driving the SR path.""" from fastvideo import VideoGenerator # noqa: PLC0415 from fastvideo.api import ( # noqa: PLC0415 - ComponentConfig, EngineConfig, GenerationRequest, GeneratorConfig, - InputConfig, OutputConfig, PipelineSelection, SamplingConfig) + ComponentConfig, EngineConfig, GenerationRequest, GeneratorConfig, InputConfig, OutputConfig, PipelineSelection, + SamplingConfig) config = GeneratorConfig( model_path=args.model, @@ -270,8 +265,7 @@ def _result_frames(result: Any) -> torch.Tensor: if frames is None and isinstance(result, dict): frames = result.get("frames") if frames is None: - raise RuntimeError( - "Generation returned no frames; ensure return_frames=True.") + raise RuntimeError("Generation returned no frames; ensure return_frames=True.") if isinstance(frames, list): frames = np.stack(frames, axis=0) if isinstance(frames, np.ndarray): @@ -309,9 +303,8 @@ def _diff(args: argparse.Namespace) -> None: ref_path = Path(args.reference) cand_path = Path(args.candidate) if not ref_path.is_file() or not cand_path.is_file(): - raise FileNotFoundError( - f"Need both --reference and --candidate to exist: {ref_path}, " - f"{cand_path}") + raise FileNotFoundError(f"Need both --reference and --candidate to exist: {ref_path}, " + f"{cand_path}") ref = torch.load(ref_path, map_location="cpu") cand = torch.load(cand_path, map_location="cpu") @@ -357,21 +350,16 @@ def _diff(args: argparse.Namespace) -> None: def _build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) sub = parser.add_subparsers(dest="mode", required=False) run = parser - run.add_argument("--label", - default="run", - help="Identifier saved alongside the output (e.g. internal/public).") - run.add_argument("--output", - default="/tmp/ltx2_sr_alignment.pt", - help="Destination .pt file.") + run.add_argument("--label", default="run", help="Identifier saved alongside the output (e.g. internal/public).") + run.add_argument("--output", default="/tmp/ltx2_sr_alignment.pt", help="Destination .pt file.") run.add_argument("--use-typed-api", action="store_true", help="Use the new typed GeneratorConfig API " - "(public side); otherwise use legacy from_pretrained.") + "(public side); otherwise use legacy from_pretrained.") run.add_argument("--model", default=DEFAULT_MODEL) run.add_argument("--seed", type=int, default=DEFAULT_SEED) run.add_argument("--height", type=int, default=DEFAULT_HEIGHT) @@ -381,9 +369,7 @@ def _build_parser() -> argparse.ArgumentParser: run.add_argument("--base-steps", type=int, default=DEFAULT_BASE_STEPS) run.add_argument("--refine-steps", type=int, default=DEFAULT_REFINE_STEPS) - run.add_argument("--diff", - action="store_true", - help="Diff two .pt outputs instead of running.") + run.add_argument("--diff", action="store_true", help="Diff two .pt outputs instead of running.") run.add_argument("--reference") run.add_argument("--candidate") return parser diff --git a/scripts/verify_anyflow_fastvideo_parity.py b/scripts/verify_anyflow_fastvideo_parity.py index 2f7b717ad..b000686c6 100644 --- a/scripts/verify_anyflow_fastvideo_parity.py +++ b/scripts/verify_anyflow_fastvideo_parity.py @@ -39,7 +39,6 @@ from pathlib import Path import torch from safetensors.torch import load_file - ANYFLOW_LOCAL = Path(os.environ.get("ANYFLOW_LOCAL", "./anyflow-1.3b")).expanduser() ANYFLOW_REF = Path(os.environ.get("ANYFLOW_REF", "./anyflow-ref")).expanduser() sys.path.insert(0, str(ANYFLOW_REF)) @@ -71,8 +70,7 @@ def init_single_rank() -> None: init_distributed_environment, initialize_model_parallel, ) - init_distributed_environment(world_size=1, rank=0, local_rank=0, - backend="nccl") + init_distributed_environment(world_size=1, rank=0, local_rank=0, backend="nccl") initialize_model_parallel( tensor_model_parallel_size=1, sequence_model_parallel_size=1, @@ -187,8 +185,8 @@ def forward_compare(fv_model, af_net): # are semantically equivalent. t_per_frame = torch.full((B, F), 500.0, device=DEVICE, dtype=DTYPE) r_per_frame = torch.full((B, F), 200.0, device=DEVICE, dtype=DTYPE) - t_per_sample = torch.full((B,), 500.0, device=DEVICE, dtype=DTYPE) - r_per_sample = torch.full((B,), 200.0, device=DEVICE, dtype=DTYPE) + t_per_sample = torch.full((B, ), 500.0, device=DEVICE, dtype=DTYPE) + r_per_sample = torch.full((B, ), 200.0, device=DEVICE, dtype=DTYPE) with torch.no_grad(): # AnyFlow native: takes [B, F, C, H, W] + [B, F] (t, r). @@ -209,8 +207,8 @@ def forward_compare(fv_model, af_net): # current timestep / attn_metadata. from fastvideo.forward_context import set_forward_context with set_forward_context( - current_timestep=t_per_sample, - attn_metadata=None, + current_timestep=t_per_sample, + attn_metadata=None, ): fv_out = fv_model( hidden_states=x, @@ -268,7 +266,8 @@ def sample_anystep(fv_model): r_timestep=r_in, ) x = scheduler.step( - v, sample=x, + v, + sample=x, timestep=t_cur.repeat(B), r_timestep=t_next.repeat(B), ) @@ -305,8 +304,8 @@ def training_step_compare(fv_model, af_net) -> bool: noise = torch.randn_like(real) t_abs_pf = torch.full((B, F), 500.0, device=DEVICE, dtype=DTYPE) r_abs_pf = torch.full((B, F), 200.0, device=DEVICE, dtype=DTYPE) - t_abs_ps = torch.full((B,), 500.0, device=DEVICE, dtype=DTYPE) - r_abs_ps = torch.full((B,), 200.0, device=DEVICE, dtype=DTYPE) + t_abs_ps = torch.full((B, ), 500.0, device=DEVICE, dtype=DTYPE) + r_abs_ps = torch.full((B, ), 200.0, device=DEVICE, dtype=DTYPE) # AnyFlow inline replica: uses [B, F, C, H, W] layout. real_btchw = real.permute(0, 2, 1, 3, 4).contiguous() @@ -315,19 +314,19 @@ def training_step_compare(fv_model, af_net) -> bool: noisy_btchw = t_norm_pf * noise_btchw + (1 - t_norm_pf) * real_btchw def u_func_af(x_in, t_in, r_in): - return af_net( - x_in, timestep=t_in, r_timestep=r_in, - encoder_hidden_states=enc, return_dict=False, is_causal=False)[0] + return af_net(x_in, + timestep=t_in, + r_timestep=r_in, + encoder_hidden_states=enc, + return_dict=False, + is_causal=False)[0] v_pred = noise_btchw - real_btchw eps = EPSILON - F_plus = u_func_af(noisy_btchw + v_pred * (eps / NUM_TRAIN_TIMESTEPS), - t_abs_pf + eps, r_abs_pf) - F_minus = u_func_af(noisy_btchw - v_pred * (eps / NUM_TRAIN_TIMESTEPS), - t_abs_pf - eps, r_abs_pf) + F_plus = u_func_af(noisy_btchw + v_pred * (eps / NUM_TRAIN_TIMESTEPS), t_abs_pf + eps, r_abs_pf) + F_minus = u_func_af(noisy_btchw - v_pred * (eps / NUM_TRAIN_TIMESTEPS), t_abs_pf - eps, r_abs_pf) dF_dt_af = (F_plus - F_minus) / (2 * eps) - target_af = ((noise_btchw - real_btchw) - - (t_abs_pf - r_abs_pf).view(B, F, 1, 1, 1) * dF_dt_af) + target_af = ((noise_btchw - real_btchw) - (t_abs_pf - r_abs_pf).view(B, F, 1, 1, 1) * dF_dt_af) flow_af = u_func_af(noisy_btchw, t_abs_pf, r_abs_pf) loss_af = (flow_af.float() - target_af.float()).pow(2).reshape(B, -1).mean(-1) @@ -347,15 +346,10 @@ def training_step_compare(fv_model, af_net) -> bool: ) v_pred_fv = noise_bcfhw - real_bcfhw - F_plus_fv = u_func_fv( - noisy_bcfhw + v_pred_fv * (eps / NUM_TRAIN_TIMESTEPS), - t_abs_ps + eps, r_abs_ps) - F_minus_fv = u_func_fv( - noisy_bcfhw - v_pred_fv * (eps / NUM_TRAIN_TIMESTEPS), - t_abs_ps - eps, r_abs_ps) + F_plus_fv = u_func_fv(noisy_bcfhw + v_pred_fv * (eps / NUM_TRAIN_TIMESTEPS), t_abs_ps + eps, r_abs_ps) + F_minus_fv = u_func_fv(noisy_bcfhw - v_pred_fv * (eps / NUM_TRAIN_TIMESTEPS), t_abs_ps - eps, r_abs_ps) dF_dt_fv = (F_plus_fv - F_minus_fv) / (2 * eps) - target_fv = ((noise_bcfhw - real_bcfhw) - - (t_abs_ps - r_abs_ps).view(B, 1, 1, 1, 1) * dF_dt_fv) + target_fv = ((noise_bcfhw - real_bcfhw) - (t_abs_ps - r_abs_ps).view(B, 1, 1, 1, 1) * dF_dt_fv) flow_fv = u_func_fv(noisy_bcfhw, t_abs_ps, r_abs_ps) loss_fv = (flow_fv.float() - target_fv.float()).pow(2).reshape(B, -1).mean(-1) @@ -392,11 +386,9 @@ def main() -> None: forward_ok = forward_compare(fv_model, af_net) sample_ok = sample_anystep(fv_model) train_ok = training_step_compare(fv_model, af_net) - banner( - f"SUMMARY: missing_keys={n_miss} unexpected_keys={n_unex} " - f"forward_parity={forward_ok} sample_smoke={sample_ok} " - f"training_parity={train_ok}" - ) + banner(f"SUMMARY: missing_keys={n_miss} unexpected_keys={n_unex} " + f"forward_parity={forward_ok} sample_smoke={sample_ok} " + f"training_parity={train_ok}") sys.exit(0 if (forward_ok and sample_ok and train_ok) else 1) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py index 70c011733..6d0f59518 100644 --- a/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py +++ b/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py @@ -30,22 +30,21 @@ PARITY_SCOPE = "both" def _tiny_config() -> DreamXWorldARConfig: - return DreamXWorldARConfig( - arch_config=DreamXWorldARArchConfig( - num_attention_heads=1, - attention_head_dim=8, - in_channels=4, - out_channels=4, - ffn_dim=16, - num_layers=1, - text_dim=8, - freq_dim=8, - text_len=4, - local_attn_size=2, - sink_size=1, - attn_compress=1, - cam_self_attn_layers=(0,), - )) + return DreamXWorldARConfig(arch_config=DreamXWorldARArchConfig( + num_attention_heads=1, + attention_head_dim=8, + in_channels=4, + out_channels=4, + ffn_dim=16, + num_layers=1, + text_dim=8, + freq_dim=8, + text_len=4, + local_attn_size=2, + sink_size=1, + attn_compress=1, + cam_self_attn_layers=(0, ), + )) def _load_official_tiny(): @@ -64,8 +63,10 @@ def _load_official_tiny(): def _sdpa_same_dtype(q, k, v, **kwargs): del kwargs - out = torch.nn.functional.scaled_dot_product_attention( - q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=0.0) + out = torch.nn.functional.scaled_dot_product_attention(q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + dropout_p=0.0) return out.transpose(1, 2).contiguous() official_attention.attention = _sdpa_same_dtype @@ -91,7 +92,7 @@ def _load_official_tiny(): add_control_adapter=True, cam_method="prope", attn_compress=1, - cam_self_attn_layers=(0,), + cam_self_attn_layers=(0, ), ).eval() @@ -136,11 +137,21 @@ def test_dreamx_world_ar_tiny_forward_matches_official(): fastvideo.load_state_dict(official.state_dict(), strict=True) x, t, context, camera, kv_cache, cross_cache = _make_inputs() - official_out = official(x=x, t=t, context=context, seq_len=4, y_camera=camera, - kv_cache=kv_cache, crossattn_cache=cross_cache).detach() + official_out = official(x=x, + t=t, + context=context, + seq_len=4, + y_camera=camera, + kv_cache=kv_cache, + crossattn_cache=cross_cache).detach() x, t, context, camera, kv_cache, cross_cache = _make_inputs() - fastvideo_out = fastvideo(x=x, t=t, context=context, seq_len=4, y_camera=camera, - kv_cache=kv_cache, crossattn_cache=cross_cache).detach() + fastvideo_out = fastvideo(x=x, + t=t, + context=context, + seq_len=4, + y_camera=camera, + kv_cache=kv_cache, + crossattn_cache=cross_cache).detach() assert official_out.abs().max() > 0, "official output is all-zero; parity comparison is vacuous" assert_close(fastvideo_out, official_out, atol=1e-5, rtol=1e-5) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py index 2448e81e2..09b9e5c2a 100644 --- a/tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py +++ b/tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py @@ -26,16 +26,14 @@ def _load_official_functions(): import importlib.util pose_path = OFFICIAL_REF_DIR / "utils" / "pose_utils.py" - pose_spec = importlib.util.spec_from_file_location( - "dreamx_world_pose_utils", pose_path) + pose_spec = importlib.util.spec_from_file_location("dreamx_world_pose_utils", pose_path) if pose_spec is None or pose_spec.loader is None: raise RuntimeError(f"Cannot load DreamX pose_utils: {pose_path}") pose_module = importlib.util.module_from_spec(pose_spec) pose_spec.loader.exec_module(pose_module) source = (OFFICIAL_REF_DIR / "utils" / "inference_utils.py").read_text() - source = source.replace( - "from .pose_utils import interpolate_camera_poses\n", "") + source = source.replace("from .pose_utils import interpolate_camera_poses\n", "") namespace = {"interpolate_camera_poses": pose_module.interpolate_camera_poses} exec(compile(source, str(OFFICIAL_REF_DIR / "utils" / "inference_utils.py"), "exec"), namespace) except Exception as exc: # noqa: BLE001 - local parity should skip missing reference deps. @@ -91,5 +89,7 @@ def test_dreamx_world_camera_conditioning_matches_official(action_seq, action_sp for key in ("viewmats", "K"): assert official[key].shape == fastvideo[key].shape diff = (official[key] - fastvideo[key]).abs() - print(f"{key}: shape={tuple(fastvideo[key].shape)} diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}") + print( + f"{key}: shape={tuple(fastvideo[key].shape)} diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}" + ) assert_close(fastvideo[key], official[key], atol=1e-5, rtol=1e-5) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py b/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py index 385c1f613..5ab0f66b2 100644 --- a/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py +++ b/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py @@ -51,7 +51,7 @@ def test_dreamx_world_5b_cam_pipeline_config_wires_first_scope_components(): assert text_config.prefix == "umt5" assert text_config.vocab_size == 256384 assert text_config.d_model == 4096 - assert config.text_encoder_precisions == ("bf16",) + assert config.text_encoder_precisions == ("bf16", ) def test_dreamx_world_pipeline_registry_discovers_entrypoint(): @@ -60,33 +60,32 @@ def test_dreamx_world_pipeline_registry_discovers_entrypoint(): assert pipelines["basic"]["DreamXWorldPipeline"] is DreamXWorldPipeline - - def test_dreamx_world_local_model_index_resolves_model_info(tmp_path): model_dir = tmp_path / "DreamX-World-5B-Cam-converted" model_dir.mkdir() for component in ("scheduler", "text_encoder", "tokenizer", "transformer", "vae"): (model_dir / component).mkdir() - (model_dir / "model_index.json").write_text(json.dumps({ - "_class_name": "DreamXWorldPipeline", - "_diffusers_version": "0.31.0", - "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], - "text_encoder": ["transformers", "UMT5EncoderModel"], - "tokenizer": ["transformers", "AutoTokenizer"], - "transformer": ["diffusers", "DreamXWorldTransformer3DModel"], - "vae": ["diffusers", "AutoencoderKLWan"], - }) + "\n") + (model_dir / "model_index.json").write_text( + json.dumps({ + "_class_name": "DreamXWorldPipeline", + "_diffusers_version": "0.31.0", + "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], + "text_encoder": ["transformers", "UMT5EncoderModel"], + "tokenizer": ["transformers", "AutoTokenizer"], + "transformer": ["diffusers", "DreamXWorldTransformer3DModel"], + "vae": ["diffusers", "AutoencoderKLWan"], + }) + "\n") info = get_model_info(str(model_dir), pipeline_type=PipelineType.BASIC, workload_type=WorkloadType.I2V) assert info.pipeline_cls is DreamXWorldPipeline assert info.pipeline_config_cls is DreamXWorld5BCamPipelineConfig + def test_dreamx_world_model_path_resolves_pipeline_config(): assert get_pipeline_config_cls_from_name("GD-ML/DreamX-World-5B-Cam") is DreamXWorld5BCamPipelineConfig - def test_dreamx_world_default_preset_is_registered(): preset_name = get_default_preset("GD-ML/DreamX-World-5B-Cam") preset = get_preset(preset_name, "dreamx_world") @@ -100,8 +99,6 @@ def test_dreamx_world_default_preset_is_registered(): assert preset.defaults["guidance_scale"] == 5.0 - - def test_dreamx_world_pipeline_initializes_official_flow_scheduler(): pipeline = DreamXWorldPipeline.__new__(DreamXWorldPipeline) pipeline.modules = {} @@ -113,6 +110,7 @@ def test_dreamx_world_pipeline_initializes_official_flow_scheduler(): assert isinstance(scheduler, FlowMatchEulerDiscreteScheduler) assert scheduler.config.shift == 3.0 + def test_dreamx_world_camera_conditioning_stage_sets_y_camera_extra(): batch = ForwardBatch( data_type="t2v", @@ -158,7 +156,6 @@ def test_dreamx_world_denoising_kwargs_filter_for_y_camera(): assert stage.prepare_extra_func_kwargs(no_y_camera, {"y_camera": y_camera}) == {} - def test_dreamx_world_ar_pipeline_config_wires_components(): config = DreamXWorld5BARPipelineConfig() assert config.is_causal is True @@ -221,7 +218,8 @@ def test_dreamx_world_camera_interpolation_handles_single_camera(): def test_dreamx_world_ar_cache_initializes_camera_self_attention_entries(): cam_self_attn = SimpleNamespace(num_heads=3, head_dim=5) transformer = SimpleNamespace( - blocks=[SimpleNamespace(cam_self_attn=cam_self_attn), SimpleNamespace(cam_self_attn=cam_self_attn)], + blocks=[SimpleNamespace(cam_self_attn=cam_self_attn), + SimpleNamespace(cam_self_attn=cam_self_attn)], num_attention_heads=2, attention_head_dim=4, ) @@ -250,9 +248,10 @@ def test_dreamx_world_ar_context_noise_fraction_maps_to_scheduler_timestep(): assert DreamXWorldARCausalDenoisingStage._context_noise_timestep(100) == 100 - def test_dreamx_world_ar_context_update_advances_camera_cache_indices(): + class DummyTransformer: + def __call__(self, *, hidden_states, encoder_hidden_states, timestep, y_camera, kv_cache, crossattn_cache, current_start): del encoder_hidden_states, y_camera, crossattn_cache @@ -290,7 +289,10 @@ def test_dreamx_world_ar_context_update_advances_camera_cache_indices(): stage._update_context_cache( block_latents=torch.zeros(1, 4, 3, 2, 2), context=[torch.zeros(2, 4)], - camera_block={"viewmats": torch.eye(4).reshape(1, 1, 4, 4), "K": torch.eye(3).reshape(1, 1, 3, 3)}, + camera_block={ + "viewmats": torch.eye(4).reshape(1, 1, 4, 4), + "K": torch.eye(3).reshape(1, 1, 3, 3) + }, kv_cache=caches, crossattn_cache=[{}], start=0, diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py index b394bbae7..367542995 100644 --- a/tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py +++ b/tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py @@ -17,8 +17,7 @@ from omegaconf import OmegaConf from torch.testing import assert_close from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler as FastVideoFlowMatchEulerDiscreteScheduler, -) + FlowMatchEulerDiscreteScheduler as FastVideoFlowMatchEulerDiscreteScheduler, ) REPO_ROOT = Path(__file__).resolve().parents[3] PARITY_SCOPE = "implementation_subcomponent" diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py index d51d83314..9b3147581 100644 --- a/tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py +++ b/tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py @@ -28,7 +28,8 @@ from fastvideo.models.loader.component_loader import TextEncoderLoader REPO_ROOT = Path(__file__).resolve().parents[3] OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World")) WAN_BASE_DIR = Path(os.getenv("DREAMX_WORLD_WAN_BASE_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B")) -WAN_DIFFUSERS_DIR = Path(os.getenv("DREAMX_WORLD_WAN_DIFFUSERS_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B-Diffusers")) +WAN_DIFFUSERS_DIR = Path( + os.getenv("DREAMX_WORLD_WAN_DIFFUSERS_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B-Diffusers")) PARITY_SCOPE = "implementation_subcomponent" @@ -55,6 +56,7 @@ def _install_xfuser_stub() -> None: distributed.model_parallel_is_initialized = lambda: False class XFuserLongContextAttention: + def __call__(self, *args, **kwargs): raise RuntimeError("xfuser stub cannot execute attention") @@ -95,9 +97,10 @@ def _load_official_text_encoder(device, dtype): from models import WanT5EncoderModel except Exception as exc: # noqa: BLE001 pytest.skip(f"Cannot import official DreamX text encoder: {exc}") - model = WanT5EncoderModel.from_pretrained( - str(text_path), additional_kwargs=_text_kwargs(), low_cpu_mem_usage=True, torch_dtype=dtype - ) + model = WanT5EncoderModel.from_pretrained(str(text_path), + additional_kwargs=_text_kwargs(), + low_cpu_mem_usage=True, + torch_dtype=dtype) return model.to(device=device, dtype=dtype).eval() diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py index fff655317..f7908fb1b 100644 --- a/tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py +++ b/tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py @@ -20,14 +20,12 @@ from omegaconf import OmegaConf from torch.testing import assert_close from fastvideo.forward_context import set_forward_context -from fastvideo.configs.models.dits.dreamx_world import ( - DreamXWorldArchConfig, DreamXWorldConfig) +from fastvideo.configs.models.dits.dreamx_world import (DreamXWorldArchConfig, DreamXWorldConfig) from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import build_dreamx_camera_condition from fastvideo.configs.pipelines.dreamx_world import make_dreamx_world_5b_cam_dit_config from fastvideo.platforms import AttentionBackendEnum -from fastvideo.models.dits.dreamx_world import ( - DreamXPropeSelfAttention, DreamXWorldTransformer3DModel, - DreamXWorldTransformerBlock) +from fastvideo.models.dits.dreamx_world import (DreamXPropeSelfAttention, DreamXWorldTransformer3DModel, + DreamXWorldTransformerBlock) from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict from fastvideo.models.loader.utils import get_param_names_mapping from fastvideo.models.loader.weight_utils import resolve_safetensors_files, safetensors_weights_iterator @@ -37,7 +35,8 @@ REPO_ROOT = Path(__file__).resolve().parents[3] OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World")) LOCAL_WEIGHTS_DIR = Path(os.getenv("DREAMX_WORLD_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / "dreamx_world")) WAN_BASE_DIR = Path(os.getenv("DREAMX_WORLD_WAN_BASE_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B")) -CONVERTED_WEIGHTS_DIR = Path(os.getenv("DREAMX_WORLD_CONVERTED_WEIGHTS_DIR", REPO_ROOT / "converted_weights" / "dreamx_world")) +CONVERTED_WEIGHTS_DIR = Path( + os.getenv("DREAMX_WORLD_CONVERTED_WEIGHTS_DIR", REPO_ROOT / "converted_weights" / "dreamx_world")) PARITY_SCOPE = "both" @@ -57,6 +56,7 @@ def _install_xfuser_stub() -> None: distributed.model_parallel_is_initialized = lambda: False class XFuserLongContextAttention: + def __call__(self, *args, **kwargs): raise RuntimeError("xfuser stub cannot execute attention") @@ -70,21 +70,20 @@ def _install_xfuser_stub() -> None: def _make_tiny_dreamx_config() -> DreamXWorldConfig: - return DreamXWorldConfig( - arch_config=DreamXWorldArchConfig( - num_attention_heads=1, - attention_head_dim=8, - in_channels=16, - out_channels=16, - ffn_dim=32, - num_layers=1, - cross_attn_norm=True, - qk_norm="rms_norm_across_heads", - add_control_adapter=True, - cam_method="prope", - attn_compress=1, - cam_self_attn_layers=None, - )) + return DreamXWorldConfig(arch_config=DreamXWorldArchConfig( + num_attention_heads=1, + attention_head_dim=8, + in_channels=16, + out_channels=16, + ffn_dim=32, + num_layers=1, + cross_attn_norm=True, + qk_norm="rms_norm_across_heads", + add_control_adapter=True, + cam_method="prope", + attn_compress=1, + cam_self_attn_layers=None, + )) def _add_official_to_path() -> None: @@ -168,9 +167,12 @@ def _make_inputs(device: torch.device, dtype: torch.dtype): context = [torch.randn(16, 4096, device=device, dtype=dtype)] seq_len = math.ceil((latent_h * latent_w) / 4 * latent_frames) timestep = torch.full((1, seq_len), 250, device=device, dtype=torch.long) - camera = build_dreamx_camera_condition( - ["w"], [4], num_frames=num_frames, height=height, width=width, dtype=dtype, device=device - ) + camera = build_dreamx_camera_condition(["w"], [4], + num_frames=num_frames, + height=height, + width=width, + dtype=dtype, + device=device) camera = {key: value.unsqueeze(0) for key, value in camera.items()} return {"x": [x[0]], "context": context, "t": timestep, "seq_len": seq_len, "y_camera": camera} @@ -186,9 +188,8 @@ def _run_official(model: torch.nn.Module, inputs: dict) -> torch.Tensor: def _run_fastvideo(model: torch.nn.Module, inputs: dict) -> torch.Tensor: hidden_states = torch.stack(inputs["x"], dim=0) - encoder_hidden_states = torch.stack([ - torch.cat([inputs["context"][0], inputs["context"][0].new_zeros(512 - inputs["context"][0].shape[0], 4096)]) - ]) + encoder_hidden_states = torch.stack( + [torch.cat([inputs["context"][0], inputs["context"][0].new_zeros(512 - inputs["context"][0].shape[0], 4096)])]) with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): output = model( hidden_states=hidden_states, @@ -217,10 +218,7 @@ def test_dreamx_world_conversion_mapping_strict_load_smoke(monkeypatch): cam_method="prope", ) official_state = official.state_dict() - diffusers_like_state = { - map_transformer_key(key): value.detach().clone() - for key, value in official_state.items() - } + diffusers_like_state = {map_transformer_key(key): value.detach().clone() for key, value in official_state.items()} import fastvideo.models.dits.dreamx_world as fastvideo_dreamx monkeypatch.setattr(fastvideo_dreamx, "get_sp_world_size", lambda: 1) @@ -268,13 +266,10 @@ def test_dreamx_world_fastvideo_prope_branch_smoke(): cam_method="prope", attn_compress=1, layer_idx=0, - supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA,), + supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA, ), ) assert block.cam_self_attn is not None - assert [ - name for name, _ in block.named_parameters() - if name.startswith("cam_self_attn.") - ][:8] == [ + assert [name for name, _ in block.named_parameters() if name.startswith("cam_self_attn.")][:8] == [ "cam_self_attn.q_proj.weight", "cam_self_attn.q_proj.bias", "cam_self_attn.k_proj.weight", @@ -313,11 +308,15 @@ def test_dreamx_world_transformer_parity_scaffold(monkeypatch): import fastvideo.models.dits.dreamx_world as fastvideo_dreamx monkeypatch.setattr(attention_layer, "get_sp_parallel_rank", lambda: 0) monkeypatch.setattr(attention_layer, "get_sp_world_size", lambda: 1) - monkeypatch.setattr(attention_layer, "sequence_model_parallel_all_to_all_4D", lambda tensor, scatter_dim=2, gather_dim=1: tensor) + monkeypatch.setattr(attention_layer, + "sequence_model_parallel_all_to_all_4D", + lambda tensor, scatter_dim=2, gather_dim=1: tensor) monkeypatch.setattr(attention_layer, "sequence_model_parallel_all_gather", lambda tensor, dim=-1: tensor) monkeypatch.setattr(fastvideo_dreamx, "get_sp_world_size", lambda: 1) monkeypatch.setattr(communication_op, "get_sp_world_size", lambda: 1) - monkeypatch.setattr(fastvideo_dreamx, "sequence_model_parallel_shard", lambda tensor, dim=1: (tensor, tensor.shape[dim])) + monkeypatch.setattr(fastvideo_dreamx, + "sequence_model_parallel_shard", + lambda tensor, dim=1: (tensor, tensor.shape[dim])) monkeypatch.setattr( fastvideo_dreamx, "sequence_model_parallel_all_gather_with_unpad", diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py index c994b4c12..ba9082386 100644 --- a/tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py +++ b/tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py @@ -50,6 +50,7 @@ def _install_xfuser_stub() -> None: distributed.model_parallel_is_initialized = lambda: False class XFuserLongContextAttention: + def __call__(self, *args, **kwargs): raise RuntimeError("xfuser stub cannot execute attention") diff --git a/tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py b/tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py index ee2af1efe..46c405a86 100644 --- a/tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py +++ b/tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py @@ -25,7 +25,6 @@ from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3 from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel from fastvideo.models.loader.component_loader import TextEncoderLoader - PARITY_SCOPE = "production_loader" MINIMAX_H3_TEXT_ENCODER_LAYER = 50 @@ -114,13 +113,10 @@ def _make_cases(root: Path) -> dict[str, dict[str, torch.Tensor]]: image_grid_thw = image_inputs["image_grid_thw"] merge_area = int(processor.image_processor.merge_size)**2 image_token_count = int(image_grid_thw[0].prod().item()) // merge_area - image_ids = ( - tokenizer(": ", add_special_tokens=False)["input_ids"] - + [_token_id(tokenizer, "<|vision_start|>")] - + [_token_id(tokenizer, "<|image_pad|>")] * image_token_count - + [_token_id(tokenizer, "<|vision_end|>")] - + prompt_ids - ) + image_ids = (tokenizer(": ", add_special_tokens=False)["input_ids"] + + [_token_id(tokenizer, "<|vision_start|>")] + + [_token_id(tokenizer, "<|image_pad|>")] * image_token_count + + [_token_id(tokenizer, "<|vision_end|>")] + prompt_ids) cases["image"] = { "input_ids": torch.tensor([image_ids], dtype=torch.long), "attention_mask": torch.ones(1, len(image_ids), dtype=torch.long), @@ -136,11 +132,9 @@ def _make_cases(root: Path) -> dict[str, dict[str, torch.Tensor]]: video_tokens_per_block = grid_height * grid_width // int(processor.video_processor.merge_size)**2 video_ids = tokenizer("