Format z_image_pipeline.py with pyink

Author: csgoogle

src/maxdiffusion/pipelines/z_image/z_image_pipeline.py

@@ -137,11 +137,7 @@ def encode_prompt(self, prompt: Union[str, list[str]], max_sequence_length: int
 
     if self.offload_encoders:
       # Restore parameters from CPU back to original device shardings
-      state = jax.tree_util.tree_map(
-          lambda x, sharding: jax.device_put(x, sharding),
-          state,
-          self._text_encoder_shardings
-      )
+      state = jax.tree_util.tree_map(lambda x, sharding: jax.device_put(x, sharding), state, self._text_encoder_shardings)
 
     with self.mesh if self.mesh is not None else nullcontext():
       embeddings = encode_step(
@@ -155,6 +151,7 @@ def encode_prompt(self, prompt: Union[str, list[str]], max_sequence_length: int
 
     if self.offload_encoders:
       from maxdiffusion import max_logging
+
       s_offload = time.perf_counter()
       cpus = jax.devices("cpu")
       offloaded_state = jax.tree_util.tree_map(lambda x: jax.device_put(x, device=cpus[0]), state)