add upload video

Author: susanbao

src/maxdiffusion/generate_wan.py

@@ -20,6 +20,50 @@
 from absl import app
 from maxdiffusion.utils import export_to_video
 import os
+from google.cloud import storage
+
+def upload_video_to_gcs(output_dir: str, video_path: str):
+    """
+    Uploads a local video file to a specified Google Cloud Storage bucket.
+    """
+    try:
+        path_without_scheme = output_dir.removeprefix("gs://")
+        parts = path_without_scheme.split('/', 1)
+        bucket_name = parts[0]
+        folder_name = parts[1] if len(parts) > 1 else ''
+
+        # Initialize the GCS client
+        storage_client = storage.Client()
+
+        # Get the bucket object
+        bucket = storage_client.bucket(bucket_name)
+
+        # Define the source and destination paths
+        source_file_path = f"./{video_path}"
+        destination_blob_name = os.path.join(folder_name, "videos", video_path)
+
+        # Create a blob object
+        blob = bucket.blob(destination_blob_name)
+
+        # Upload the file
+        print(f"Uploading {source_file_path} to {bucket_name}/{destination_blob_name}...")
+        blob.upload_from_filename(source_file_path)
+        print(f"Upload complete {source_file_path}.")
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+def delete_file(file_path: str):
+   # Best practice: Check if the file exists before trying to delete it.
+  if os.path.exists(file_path):
+      try:
+          os.remove(file_path)
+          print(f"Successfully deleted file: {file_path}")
+      except OSError as e:
+          # This catches other issues like permission errors
+          print(f"Error deleting file '{file_path}': {e}")
+  else:
+      print(f"The file '{file_path}' does not exist.")
 
 jax.config.update("jax_use_shardy_partitioner", True)
 
@@ -44,8 +88,10 @@ def inference_generate_video(config, pipeline, filename_prefix=""):
 
   print(f"video {filename_prefix}, compile time: {(time.perf_counter() - s0)}")
   for i in range(len(videos)):
-    video_path = os.path.join(config.output_dir, "videos", f"{filename_prefix}wan_output_{config.seed}_{i}.mp4")
+    video_path = f"{filename_prefix}wan_output_{config.seed}_{i}.mp4"
     export_to_video(videos[i], video_path, fps=config.fps)
+    upload_video_to_gcs(config.output_dir, video_path)
+    delete_file(f"./{video_path}")
   return
 
 def run(config, pipeline=None, filename_prefix=""):
@@ -79,9 +125,10 @@ def run(config, pipeline=None, filename_prefix=""):
   print("compile time: ", (time.perf_counter() - s0))
   saved_video_path = []
   for i in range(len(videos)):
-    video_path = os.path.join(config.output_dir, "videos", f"{filename_prefix}wan_output_{config.seed}_{i}.mp4")
+    video_path = os.path.join(f"{filename_prefix}wan_output_{config.seed}_{i}.mp4")
     export_to_video(videos[i], video_path, fps=config.fps)
     saved_video_path.append(video_path)
+    upload_video_to_gcs(config.output_dir, video_path)
 
   s0 = time.perf_counter()
   videos = pipeline(

src/maxdiffusion/trainers/wan_trainer.py

@@ -153,8 +153,8 @@ def start_training(self):
     pretrained_video_path = generate_sample(self.config, pipeline, filename_prefix="pre-training-")
 
     # save some memory.
-    del pipeline.vae
-    del pipeline.vae_cache
+    # del pipeline.vae
+    # del pipeline.vae_cache
 
     mesh = pipeline.mesh
     train_data_iterator = self.load_dataset(mesh, is_training=True)
@@ -252,6 +252,7 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
         if self.config.eval_every > 0 and (step + 1) % self.config.eval_every == 0:
           # Re-create the iterator each time you start evaluation to reset it
           # This assumes your data loading logic can be called to get a fresh iterator.
+          pipeline.transformer = nnx.merge(state.graphdef, state.params, state.rest_of_state)
           inference_generate_video(self.config, pipeline, filename_prefix=f"{step+1}-train_steps-")
           eval_data_iterator = self.load_dataset(mesh, is_training=False)
           eval_rng = jax.random.key(self.config.seed + step)