FLUX: add load_dataset to __ini__ file

flux/flux/__init__.py

@@ -1,7 +1,5 @@
 # Copyright © 2024 Apple Inc.
 
-import math
-import time
 from typing import Tuple
 
 import mlx.core as mx
@@ -9,6 +7,7 @@ import mlx.nn as nn
 from mlx.utils import tree_unflatten
 from tqdm import tqdm
 
+from .datasets import load_dataset
 from .lora import LoRALinear
 from .sampler import FluxSampler
 from .utils import (