enable vae slicing

Author: prishajain1

src/maxdiffusion/generate_ltx2.py

@@ -119,6 +119,9 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
     checkpoint_loader = LTX2Checkpointer(config=config)
     pipeline, _, _ = checkpoint_loader.load_checkpoint()
 
+  pipeline.enable_vae_slicing()
+  pipeline.enable_vae_tiling()
+
   s0 = time.perf_counter()
 
   # Using global_batch_size_to_train_on to map prompts

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -260,6 +260,20 @@ def _init_dummy_shape(node):
           return jnp.zeros(node.shape, dtype=node.dtype)
       return node
 
+  def enable_vae_slicing(self):
+      self.vae.use_slicing = True
+
+  def disable_vae_slicing(self):
+      self.vae.use_slicing = False
+
+  def enable_vae_tiling(self):
+      if hasattr(self.vae, "enable_tiling"):
+          self.vae.enable_tiling()
+      self.vae.use_tiling = True
+
+  def disable_vae_tiling(self):
+      self.vae.use_tiling = False
+
   @classmethod
   def load_tokenizer(cls, config: HyperParameters):
       max_logging.log("Loading Gemma Tokenizer...")