fix(flux2klein): use dynamic joint_attention_dim in AOT graph compilation for 4B and 9B

Author: amepas

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -119,7 +119,7 @@ def compile_aot_async(self, params, vae_params, qwen3_params, batch_size=1, heig
 
     dummy_latents = jnp.zeros((batch_size, seq_len_img, 128), dtype=jnp.float32)
     dummy_img_ids = jnp.zeros((batch_size, seq_len_img, 4), dtype=jnp.int32)
-    dummy_prompt_embeds = jnp.zeros((batch_size, seq_len_txt, 12288), dtype=jnp.bfloat16)
+    dummy_prompt_embeds = jnp.zeros((batch_size, seq_len_txt, self.transformer.joint_attention_dim), dtype=jnp.bfloat16)
     dummy_txt_ids = jnp.zeros((batch_size, seq_len_txt, 4), dtype=jnp.float32)
     dummy_t_vec = jnp.zeros((batch_size,), dtype=jnp.float32)