Errors catching corrections

Author: prishajain1

src/maxdiffusion/generate_ltx2.py

@@ -22,6 +22,7 @@
 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.utils.export_utils import export_to_video_with_audio
 
@@ -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):

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -4,6 +4,7 @@
 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 +102,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)