fix

Author: prishajain1

src/maxdiffusion/generate_ltx2.py

@@ -22,8 +22,9 @@
 from maxdiffusion import pyconfig, max_logging, max_utils
 from absl import app
 from google.cloud import storage
+from google.api_core.exceptions import GoogleAPIError
 import flax
-from maxdiffusion.pipelines.ltx2.ltx2_pipeline_utils import encode_video
+from maxdiffusion.utils.export_utils import export_to_video_with_audio
 
 
 def upload_video_to_gcs(output_dir: str, video_path: str):
@@ -48,8 +49,8 @@ def upload_video_to_gcs(output_dir: str, video_path: str):
     blob.upload_from_filename(source_file_path)
     max_logging.log(f"Upload complete {source_file_path}.")
 
-  except Exception as e:
-    max_logging.log(f"An error occurred: {e}")
+  except GoogleAPIError as e:
+    max_logging.log(f"A storage error occurred during upload: {e}")
 
 
 def delete_file(file_path: str):
@@ -163,7 +164,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
     video_path = f"{filename_prefix}ltx2_output_{getattr(config, 'seed', 0)}_{i}.mp4"
     audio_i = audios[i] if audios is not None else None
 
-    encode_video(video=videos[i], fps=fps, audio=audio_i, audio_sample_rate=audio_sample_rate, output_path=video_path)
+    export_to_video_with_audio(video=videos[i], fps=fps, audio=audio_i, audio_sample_rate=audio_sample_rate, output_path=video_path)
 
     saved_video_path.append(video_path)
     if config.output_dir.startswith("gs://"):

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -355,13 +355,13 @@ def __init__(
     # 1. Define Partitioned Initializers (Logical Axes)
     # Q, K, V kernels: [in_features (embed), out_features (heads)]
     qkv_kernel_init = nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "heads"))
-    # Q, K, V biases: [out_features (heads)]
-    qkv_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), ("heads",))
+    # Q, K, V biases: [out_features (embed)]
+    qkv_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), ("embed",))
 
     # Out kernel: [in_features (heads), out_features (embed)]
     out_kernel_init = nnx.with_partitioning(nnx.initializers.lecun_normal(), ("heads", "embed"))
-    # Out bias: [out_features (embed)]
-    out_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), ("embed",))
+    # Out bias: [out_features (heads)]
+    out_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), ("heads",))
 
     # Norm scales
     norm_scale_init = nnx.with_partitioning(nnx.initializers.ones_init(), ("norm",))

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -1,9 +1,26 @@
+"""
+Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
 import json
 import torch
 import jax
 import jax.numpy as jnp
 from maxdiffusion import max_logging
 from huggingface_hub import hf_hub_download
+from huggingface_hub.utils import EntryNotFoundError
 from safetensors import safe_open
 from flax.traverse_util import unflatten_dict, flatten_dict
 from ..modeling_flax_pytorch_utils import (rename_key, rename_key_and_reshape_tensor, torch2jax, validate_flax_state_dict)
@@ -101,12 +118,12 @@ def load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device):
       with safe_open(shard_path, framework="pt") as f:
         for k in f.keys():
           tensors[k] = torch2jax(f.get_tensor(k))
-  except Exception:
+  except EntryNotFoundError:
     # Fallback to single file
     filename = "diffusion_pytorch_model.safetensors"
     try:
       ckpt_path = hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=filename)
-    except Exception:
+    except EntryNotFoundError:
       filename = "diffusion_pytorch_model.bin"
       ckpt_path = hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=filename)
 

src/maxdiffusion/models/ltx2/text_encoders/embeddings_connector_ltx2.py

@@ -1,5 +1,5 @@
 """
-Copyright 2025 Google LLC
+Copyright 2026 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.

src/maxdiffusion/models/ltx2/text_encoders/feature_extractor_ltx2.py

@@ -1,5 +1,5 @@
 """
-Copyright 2025 Google LLC
+Copyright 2026 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.

src/maxdiffusion/models/ltx2/text_encoders/text_encoders_ltx2.py

@@ -1,5 +1,5 @@
 """
-Copyright 2025 Google LLC
+Copyright 2026 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -1,5 +1,5 @@
 """
-Copyright 2025 Google LLC
+Copyright 2026 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -837,8 +837,8 @@ def init_block(rngs):
         rngs=rngs,
         dtype=self.dtype,
         param_dtype=self.weights_dtype,
-        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, "embed")),
-        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
+        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("embed", None)),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, (None,)),
     )
 
     self.audio_norm_out = nnx.LayerNorm(
@@ -850,8 +850,8 @@ def init_block(rngs):
         rngs=rngs,
         dtype=self.dtype,
         param_dtype=self.weights_dtype,
-        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, "embed")),
-        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
+        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("embed", None)),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, (None,)),
     )
 
   def __call__(

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline_utils.py

@@ -1,5 +1,5 @@
 """
-Copyright 2025 Google LLC
+Copyright 2026 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.

src/maxdiffusion/tests/ltx2/test_checkpointer_ltx2.py

@@ -1,5 +1,5 @@
 """
-Copyright 2025 Google LLC
+Copyright 2026 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.