diff --git a/src/flux_1/flux.py b/src/flux_1/flux.py index 8ffbab7..d1eacb3 100644 --- a/src/flux_1/flux.py +++ b/src/flux_1/flux.py @@ -4,9 +4,8 @@ from PIL import Image from tqdm import tqdm from flux_1.config.config import Config -from flux_1.config.runtime_config import RuntimeConfig -from flux_1.latent_creator.latent_creator import LatentCreator from flux_1.config.model_config import ModelConfig +from flux_1.config.runtime_config import RuntimeConfig from flux_1.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder from flux_1.models.text_encoder.t5_encoder.t5_encoder import T5Encoder from flux_1.models.transformer.transformer import Transformer @@ -48,7 +47,10 @@ class Flux1: config = RuntimeConfig(config, self.model_config) # Create the latents - latents = LatentCreator.create(config.height, config.width, seed) + latents = mx.random.normal( + shape=[1, (config.height // 16) * (config.width // 16), 64], + key=mx.random.key(seed) + ) # Embedd the prompt t5_tokens = self.t5_tokenizer.tokenize(prompt) diff --git a/src/flux_1/latent_creator/__init__.py b/src/flux_1/latent_creator/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/flux_1/latent_creator/latent_creator.py b/src/flux_1/latent_creator/latent_creator.py deleted file mode 100644 index aa251db..0000000 --- a/src/flux_1/latent_creator/latent_creator.py +++ /dev/null @@ -1,11 +0,0 @@ -import mlx.core as mx - - -class LatentCreator: - - @staticmethod - def create(height: int, width: int, seed: int) -> (mx.array, mx.array): - latent_height = height // 16 - latent_width = width // 16 - latents = mx.random.normal(shape=[1, latent_height * latent_width, 64], key=mx.random.key(seed)) - return latents