Disable unsafe rng

Author: eltsai

src/maxdiffusion/generate_wan.py

@@ -76,15 +76,6 @@ def get_git_commit_hash():
     return None
 
 jax.config.update("jax_use_shardy_partitioner", True)
-jax.config.update("jax_default_prng_impl", "unsafe_rbg")
-  # TF allocates extraneous GPU memory when using TFDS data
-  # this leads to CUDA OOMs. WAR for now is to hide GPUs from TF
-  # tf.config.set_visible_devices([], "GPU")
-if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""):
-  max_logging.log("Enabling unsafe RNG bit generator for TPU SPMD.")
-  os.environ["LIBTPU_INIT_ARGS"] = (
-      os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
-  )
 
 def call_pipeline(config, pipeline, prompt, negative_prompt):
   model_key = config.model_name