AQT integration

[mailvijayasingh]

No description provided.

[mailvijayasingh]

Verified generate sdxl offline on v6e-4 and images look fine.

[shungcp]

Hello Vijaya,

I met this error message when I follow up the maxdiffusion on v6e cloud document:

kubectl logs -l app=max-diffusion-server
     ^^^^^^
  File "/usr/local/lib/python3.12/site-packages/maxdiffusion/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py", line 261, in _generate
    latents, _ = jax.lax.fori_loop(0, num_inference_steps, loop_body, (latents, scheduler_state))
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/maxdiffusion/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py", line 242, in loop_body
    {"params": params["unet"], "aqt": params["unet"]["aqt"]},
                                      ~~~~~~~~~~~~~~^^^^^^^
KeyError: 'aqt'
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

Can you help verify if it works on v6e-1 on gke?