fixes

Author: sayakpaul

benchmarks/benchmarking_flux.py

@@ -1,17 +1,20 @@
 import torch
+from benchmarking_utils import BenchmarkMixin
 
 from diffusers import FluxTransformer2DModel
 from diffusers.utils.testing_utils import torch_device
 
-from .benchmarking_utils import BenchmarkMixin
-
 
 class BenchmarkFlux(BenchmarkMixin):
     model_class = FluxTransformer2DModel
     compile_kwargs = {"fullgraph": True, "mode": "max-autotune"}
 
     def get_model_init_dict(self):
-        return {"ckpt_id": "black-forest-labs/FLUX.1-dev", "subfolder": "transformer", "torch_dtype": torch.bfloat16}
+        return {
+            "pretrained_model_name_or_path": "black-forest-labs/FLUX.1-dev",
+            "subfolder": "transformer",
+            "torch_dtype": torch.bfloat16,
+        }
 
     def initialize_model(self):
         model = self.model_class.from_pretrained(**self.get_model_init_dict())

benchmarks/benchmarking_utils.py

@@ -1,7 +1,7 @@
 import gc
 
 import torch
-from torch.utils.benchmark import benchmark
+import torch.utils.benchmark as benchmark
 
 from diffusers.models.modeling_utils import ModelMixin
 from diffusers.utils.testing_utils import require_torch_gpu