use State instead of GraphState

PiperOrigin-RevId: 897380120

Author: Cristian Garcia

src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py

@@ -689,7 +689,7 @@ def prepare_video_latents(
 @partial(jax.jit, static_argnames=("do_classifier_free_guidance", "guidance_scale"))
 def transformer_forward_pass(
     graphdef: nnx.graph.GraphDef,
-    sharded_state: nnx.graph.GraphState,
+    sharded_state: nnx.State,
     rest_of_state: Any,
     latents: jax.Array,
     timestep: jax.Array,