Add note for skipping prk

Author: skirsten

src/diffusers/schedulers/scheduling_pndm_flax.py

@@ -388,13 +388,14 @@ def step_plms(
                 "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
             )
 
-        if not self.config.skip_prk_steps:
-            raise ValueError(
-                f"{self.__class__} can only be run AFTER scheduler has been run "
-                "in 'prk' mode for at least 12 iterations "
-                "See: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pipeline_pndm.py "
-                "for more information."
-            )
+        # NOTE: There is no way to check in the jitted runtime if the prk mode was ran before
+        # if not self.config.skip_prk_steps and state.counter < 3:
+        #     raise ValueError(
+        #         f"{self.__class__} can only be run AFTER scheduler has been run "
+        #         "in 'prk' mode for at least 12 iterations "
+        #         "See: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pipeline_pndm.py "
+        #         "for more information."
+        #     )
 
         prev_timestep = timestep - self.config.num_train_timesteps // state.num_inference_steps
         prev_timestep = jnp.where(prev_timestep > 0, prev_timestep, 0)