remove unused dependencies.

Author: entrpn

src/maxdiffusion/generate_flux.py

@@ -14,7 +14,7 @@
  limitations under the License.
 """
 
-from typing import Any, Callable, Dict, List, Optional, Union, Sequence
+from typing import Callable, List, Union, Sequence
 from absl import app
 import functools
 import math

src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py

@@ -28,9 +28,7 @@
 from ...embeddings_flax import (FluxPosEmbed, CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings)
 from .... import common_types
 from ....common_types import BlockSizes
-from .... import max_logging
 from ....utils import BaseOutput
-from dataclasses import dataclass
 
 AxisNames = common_types.AxisNames
 BATCH = common_types.BATCH

src/maxdiffusion/models/flux/util.py

@@ -127,7 +127,7 @@ def validate_flax_state_dict(expected_pytree: dict, new_pytree: dict):
     if key in new_pytree.keys():
       try:
         expected_pytree_shape = expected_pytree[key].shape
-      except:
+      except Exception:
         expected_pytree_shape = expected_pytree[key].value.shape
       if expected_pytree_shape != new_pytree[key].shape:
         max_logging.log(

src/maxdiffusion/tests/text_encoders_test.py

@@ -17,7 +17,6 @@
 import os
 import unittest
 import pytest
-import jax.numpy as jnp
 from absl.testing import absltest
 
 from transformers import CLIPTokenizer, FlaxCLIPTextModel

src/maxdiffusion/tests/vae_test.py

@@ -24,7 +24,6 @@
 import jax
 import jax.numpy as jnp
 from maxdiffusion import FlaxAutoencoderKL
-from maxdiffusion.image_processor import VaeImageProcessor
 from skimage.metrics import structural_similarity as ssim
 
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))