more linter fixes.

Author: jfacevedo-google

src/maxdiffusion/models/flux/util.py

@@ -4,7 +4,6 @@
 
 import jax
 from jax.typing import DTypeLike
-import torch  # need for torch 2 jax
 from chex import Array
 from flax.traverse_util import flatten_dict, unflatten_dict
 from huggingface_hub import hf_hub_download

src/maxdiffusion/models/wan/wan_utils.py

@@ -3,7 +3,7 @@
 from maxdiffusion import max_logging
 from huggingface_hub import hf_hub_download
 from safetensors import safe_open
-from flax.traverse_util import flatten_dict, unflatten_dict
+from flax.traverse_util import unflatten_dict
 from ..modeling_flax_pytorch_utils import (rename_key, rename_key_and_reshape_tensor, torch2jax, validate_flax_state_dict)
 
 
@@ -12,7 +12,7 @@ def _tuple_str_to_int(in_tuple):
   for item in in_tuple:
     try:
       out_list.append(int(item))
-    except:
+    except ValueError:
       out_list.append(item)
   return tuple(out_list)
 

src/maxdiffusion/tests/wan_vae_test.py

@@ -23,7 +23,6 @@
 from flax import nnx
 import numpy as np
 import unittest
-import pytest
 from absl.testing import absltest
 from skimage.metrics import structural_similarity as ssim
 from ..models.wan.autoencoder_kl_wan import (
@@ -172,7 +171,7 @@ def test_wanrms_norm(self):
     dummy_input = jnp.ones(input_shape)
     output = wanrms_norm(dummy_input)
     output_np = np.array(output)
-    assert np.allclose(output_np, torch_output_np) == True
+    assert np.allclose(output_np, torch_output_np) is True
 
     # --- Test Case 2: images == False ---
     model = TorchWanRMS_norm(dim, images=False)
@@ -186,7 +185,7 @@ def test_wanrms_norm(self):
     dummy_input = jnp.ones(input_shape)
     output = wanrms_norm(dummy_input)
     output_np = np.array(output)
-    assert np.allclose(output_np, torch_output_np) == True
+    assert np.allclose(output_np, torch_output_np) is True
 
   def test_zero_padded_conv(self):
 
@@ -235,8 +234,6 @@ def test_wan_resample(self):
     w = 720
     mode = "downsample2d"
     input_shape = (batch, dim, t, h, w)
-    expected_output_shape = (1, dim, 1, 240, 360)
-    # output dim should be (1, 96, 1, 480, 720)
     dummy_input = torch.ones(input_shape)
     torch_wan_resample = TorchWanResample(dim=dim, mode=mode)
     torch_output = torch_wan_resample(dummy_input)
@@ -426,7 +423,7 @@ def vae_encode(video, wan_vae, vae_cache, key):
     rngs = nnx.Rngs(key)
     wan_vae = AutoencoderKLWan.from_config(pretrained_model_name_or_path, subfolder="vae", rngs=rngs)
     vae_cache = AutoencoderKLWanCache(wan_vae)
-    video_path, fps = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/hiker.mp4", 8
+    video_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/hiker.mp4"
     video = load_video(video_path)
 
     vae_scale_factor_spatial = 2 ** len(wan_vae.temperal_downsample)

src/maxdiffusion/utils/loading_utils.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, Callable, List, Optional, Tuple, Union
+from typing import Callable, List, Optional, Union
 
 import PIL.Image
 import PIL.ImageOps