precompile for diffusion

[pathfinder-pf]

Motivation

Modifications

• fix flax version to 0.12.4, because flax==0.12.5 has breaking changes
• precompile for diffusion stage
• change server parameter from vae-decode-precompile-width-height to precompile-width-heights, --vae-decode-precompile-frame-paddings to precompile-frame-paddings

Accuracy Tests

JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 -u -m sgl_jax.launch_server --multimodal --model-path=/models/Wan-AI/Wan2.1-T2V-1.3B-Diffusers --precompile-token-paddings=8192 --precompile-bs-paddings=64 --precompile-width-heights 480*832 --precompile-frame-paddings 13 41

curl http://localhost:30000/api/v1/videos/generation -H "Content-Type: application/json" -d '{"prompt": "A curious raccoon", "size": "480*832", "num_frames": 41, "num_inference_steps": 50}'

8a47ddcaf0474c78a4934cda30b9ee7e.mp4

Benchmarking and Profiling

1st 58.729s
2nd 58.506s

Checklist

• Please use English, otherwise it will be closed.
• The purpose of the PR, or link existing issues this PR will resolve.
• The test plan, such as providing test command.
• (Optional) The necessary documentation update.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!