Refactor: Add separate Flux initializer

src/mflux/config/runtime_config.py

@@ -10,7 +10,11 @@ logger = logging.getLogger(__name__)
 
 
 class RuntimeConfig:
-    def __init__(self, config: Config | ConfigControlnet, model_config):
+    def __init__(
+        self,
+        config: Config | ConfigControlnet,
+        model_config: ModelConfig,
+    ):
         self.config = config
         self.model_config = model_config
         self.sigmas = self._create_sigmas(config, model_config)

src/mflux/controlnet/controlnet_util.py

@@ -2,15 +2,47 @@ import logging
 import os
 
 import cv2
+import mlx.core as mx
 import numpy as np
 import PIL.Image
 
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.models.vae.vae import VAE
+from mflux.post_processing.array_util import ArrayUtil
+
 log = logging.getLogger(__name__)
 
 
 class ControlnetUtil:
     @staticmethod
-    def preprocess_canny(img: PIL.Image) -> PIL.Image:
+    def encode_image(
+        vae: VAE,
+        config: RuntimeConfig,
+        controlnet_image_path: str,
+        controlnet_save_canny: bool,
+        output: str,
+    ) -> mx.array:
+        from mflux import ImageUtil
+
+        control_image = ImageUtil.load_image(controlnet_image_path)
+        control_image = ControlnetUtil._scale_image(config.height, config.width, control_image)
+        control_image = ControlnetUtil._preprocess_canny(control_image)
+
+        if controlnet_save_canny:
+            base, ext = os.path.splitext(output)
+            ImageUtil.save_image(
+                image=control_image,
+                path=f"{base}_controlnet_canny{ext}"
+            )  # fmt: off
+
+        controlnet_cond = ImageUtil.to_array(control_image)
+        controlnet_cond = vae.encode(controlnet_cond)
+        controlnet_cond = (controlnet_cond / vae.scaling_factor) + vae.shift_factor
+        controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=config.height, width=config.width)
+        return controlnet_cond
+
+    @staticmethod
+    def _preprocess_canny(img: PIL.Image) -> PIL.Image:
         image_to_canny = np.array(img)
         image_to_canny = cv2.Canny(image_to_canny, 100, 200)
         image_to_canny = np.array(image_to_canny[:, :, None])
@@ -18,16 +50,8 @@ class ControlnetUtil:
         return PIL.Image.fromarray(image_to_canny)
 
     @staticmethod
-    def scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image:
+    def _scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image:
         if height != img.height or width != img.width:
             log.warning(f"Control image has different dimensions than the model. Resizing to {width}x{height}")
             img = img.resize((width, height), PIL.Image.LANCZOS)
         return img
-
-    @staticmethod
-    def save_canny_image(control_image: PIL.Image, path: str):
-        from mflux import ImageUtil
-
-        base, ext = os.path.splitext(path)
-        new_filename = f"{base}_controlnet_canny{ext}"
-        ImageUtil.save_image(control_image, new_filename)

src/mflux/controlnet/flux_controlnet.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
 import mlx.core as mx
+from mlx import nn
 from tqdm import tqdm
 
 from mflux.config.config import ConfigControlnet
@@ -8,8 +9,8 @@ from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.controlnet.controlnet_util import ControlnetUtil
 from mflux.controlnet.transformer_controlnet import TransformerControlnet
-from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
 from mflux.error.exceptions import StopImageGenerationException
+from mflux.flux.flux_initializer import FluxInitializer
 from mflux.latent_creator.latent_creator import LatentCreator
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
 from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
@@ -19,16 +20,16 @@ from mflux.post_processing.array_util import ArrayUtil
 from mflux.post_processing.generated_image import GeneratedImage
 from mflux.post_processing.image_util import ImageUtil
 from mflux.post_processing.stepwise_handler import StepwiseHandler
-from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
-from mflux.tokenizer.t5_tokenizer import TokenizerT5
-from mflux.tokenizer.tokenizer_handler import TokenizerHandler
 from mflux.weights.model_saver import ModelSaver
-from mflux.weights.weight_handler import WeightHandler
-from mflux.weights.weight_handler_lora import WeightHandlerLoRA
-from mflux.weights.weight_util import WeightUtil
 
 
-class Flux1Controlnet:
+class Flux1Controlnet(nn.Module):
+    vae: VAE
+    transformer: Transformer
+    transformer_controlnet: TransformerControlnet
+    t5_text_encoder: T5Encoder
+    clip_text_encoder: CLIPEncoder
+
     def __init__(
         self,
         model_config: ModelConfig,
@@ -38,45 +39,14 @@ class Flux1Controlnet:
         lora_scales: list[float] | None = None,
         controlnet_path: str | None = None,
     ):
-        self.lora_paths = lora_paths
+        super().__init__()
-        self.lora_scales = lora_scales
+        FluxInitializer.init_controlnet(
-        self.model_config = model_config
+            flux_model=self,
-
+            model_config=model_config,
-        # Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface
+            quantize=quantize,
-        tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path)
+            local_path=local_path,
-        self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length)
+            lora_paths=lora_paths,
-        self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
+            lora_scales=lora_scales,
-
-        # Load the weights
-        weights = WeightHandler.load_regular_weights(repo_id=model_config.model_name, local_path=local_path)
-
-        # Initialize the models
-        self.vae = VAE()
-        self.transformer = Transformer(model_config, num_transformer_blocks=weights.num_transformer_blocks(), num_single_transformer_blocks=weights.num_single_transformer_blocks())  # fmt: off
-        self.t5_text_encoder = T5Encoder()
-        self.clip_text_encoder = CLIPEncoder()
-
-        # Set the weights and quantize the model
-        self.bits = WeightUtil.set_weights_and_quantize(
-            quantize_arg=quantize,
-            weights=weights,
-            vae=self.vae,
-            transformer=self.transformer,
-            t5_text_encoder=self.t5_text_encoder,
-            clip_text_encoder=self.clip_text_encoder,
-        )
-
-        # Set LoRA weights
-        lora_weights = WeightHandlerLoRA.load_lora_weights(transformer=self.transformer, lora_files=lora_paths, lora_scales=lora_scales)  # fmt:off
-        WeightHandlerLoRA.set_lora_weights(transformer=self.transformer, loras=lora_weights)
-
-        # Set Controlnet weights
-        weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer()
-        self.transformer_controlnet = TransformerControlnet(model_config=model_config, num_transformer_blocks=weights_controlnet.num_transformer_blocks(), num_single_transformer_blocks=weights_controlnet.num_single_transformer_blocks())  # fmt:off
-        WeightUtil.set_controlnet_weights_and_quantize(
-            quantize_arg=quantize,
-            weights=weights_controlnet,
-            transformer_controlnet=self.transformer_controlnet,
         )
 
     def generate_image(
@@ -88,7 +58,7 @@ class Flux1Controlnet:
         controlnet_save_canny: bool = False,
         config: ConfigControlnet = ConfigControlnet(),
         stepwise_output_dir: Path = None,
-    ) -> GeneratedImage:  # fmt: off
+    ) -> GeneratedImage:
         # Create a new runtime config based on the model type and input parameters
         config = RuntimeConfig(config, self.model_config)
         time_steps = tqdm(range(config.num_inference_steps))
@@ -101,11 +71,21 @@ class Flux1Controlnet:
             output_dir=stepwise_output_dir,
         )
 
-        # 0. Embed the controlnet reference image
+        # 0. Encode the controlnet reference image
-        controlnet_condition = self._embed_image(config, controlnet_image_path, controlnet_save_canny, output)
+        controlnet_condition = ControlnetUtil.encode_image(
+            vae=self.vae,
+            config=config,
+            controlnet_image_path=controlnet_image_path,
+            controlnet_save_canny=controlnet_save_canny,
+            output=output,
+        )
 
         # 1. Create the initial latents
-        latents = LatentCreator.create(seed=seed, height=config.height, width=config.width)
+        latents = LatentCreator.create(
+            seed=seed,
+            height=config.height,
+            width=config.width
+        )  # fmt: off
 
         # 2. Embed the prompt
         t5_tokens = self.t5_tokenizer.tokenize(prompt)
@@ -165,26 +145,6 @@ class Flux1Controlnet:
             controlnet_image_path=controlnet_image_path,
         )
 
-    def _embed_image(
-        self,
-        config: RuntimeConfig,
-        controlnet_image_path: str,
-        controlnet_save_canny: bool,
-        output: str,
-    ):
-        control_image = ImageUtil.load_image(controlnet_image_path)
-        control_image = ControlnetUtil.scale_image(config.height, config.width, control_image)
-        control_image = ControlnetUtil.preprocess_canny(control_image)
-
-        if controlnet_save_canny:
-            ControlnetUtil.save_canny_image(control_image, output)
-
-        controlnet_cond = ImageUtil.to_array(control_image)
-        controlnet_cond = self.vae.encode(controlnet_cond)
-        controlnet_cond = (controlnet_cond / self.vae.scaling_factor) + self.vae.shift_factor
-        controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=config.height, width=config.width)
-        return controlnet_cond
-
     def save_model(self, base_path: str) -> None:
         ModelSaver.save_model(self, self.bits, base_path)
         ModelSaver.save_weights(base_path, self.bits, self.transformer_controlnet, "transformer_controlnet")

src/mflux/flux/flux.py

@@ -9,6 +9,7 @@ from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig, ModelLookup
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.error.exceptions import StopImageGenerationException
+from mflux.flux.flux_initializer import FluxInitializer
 from mflux.latent_creator.latent_creator import LatentCreator
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
 from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
@@ -18,16 +19,15 @@ from mflux.post_processing.array_util import ArrayUtil
 from mflux.post_processing.generated_image import GeneratedImage
 from mflux.post_processing.image_util import ImageUtil
 from mflux.post_processing.stepwise_handler import StepwiseHandler
-from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
-from mflux.tokenizer.t5_tokenizer import TokenizerT5
-from mflux.tokenizer.tokenizer_handler import TokenizerHandler
 from mflux.weights.model_saver import ModelSaver
-from mflux.weights.weight_handler import WeightHandler
-from mflux.weights.weight_handler_lora import WeightHandlerLoRA
-from mflux.weights.weight_util import WeightUtil
 
 
 class Flux1(nn.Module):
+    vae: VAE
+    transformer: Transformer
+    t5_text_encoder: T5Encoder
+    clip_text_encoder: CLIPEncoder
+
     def __init__(
         self,
         model_config: ModelConfig,
@@ -37,38 +37,15 @@ class Flux1(nn.Module):
         lora_scales: list[float] | None = None,
     ):
         super().__init__()
-        self.lora_paths = lora_paths
+        FluxInitializer.init(
-        self.lora_scales = lora_scales
+            flux_model=self,
-        self.model_config = model_config
+            model_config=model_config,
-
+            quantize=quantize,
-        # Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface
+            local_path=local_path,
-        tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path)
+            lora_paths=lora_paths,
-        self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length)
+            lora_scales=lora_scales,
-        self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
-
-        # Load the weights
-        weights = WeightHandler.load_regular_weights(repo_id=model_config.model_name, local_path=local_path)
-
-        # Initialize the models
-        self.vae = VAE()
-        self.transformer = Transformer(model_config, num_transformer_blocks=weights.num_transformer_blocks(), num_single_transformer_blocks=weights.num_single_transformer_blocks())  # fmt: off
-        self.t5_text_encoder = T5Encoder()
-        self.clip_text_encoder = CLIPEncoder()
-
-        # Set the weights and quantize the model
-        self.bits = WeightUtil.set_weights_and_quantize(
-            quantize_arg=quantize,
-            weights=weights,
-            vae=self.vae,
-            transformer=self.transformer,
-            t5_text_encoder=self.t5_text_encoder,
-            clip_text_encoder=self.clip_text_encoder,
         )
 
-        # Set LoRA weights
-        lora_weights = WeightHandlerLoRA.load_lora_weights(transformer=self.transformer, lora_files=lora_paths, lora_scales=lora_scales)  # fmt:off
-        WeightHandlerLoRA.set_lora_weights(transformer=self.transformer, loras=lora_weights)
-
     def generate_image(
         self,
         seed: int,
@@ -89,7 +66,11 @@ class Flux1(nn.Module):
         )
 
         # 1. Create the initial latents
-        latents = LatentCreator.create_for_txt2img_or_img2img(seed, config, self.vae)
+        latents = LatentCreator.create_for_txt2img_or_img2img(
+            seed=seed,
+            vae=self.vae,
+            runtime_conf=config,
+        )
 
         # 2. Embed the prompt
         t5_tokens = self.t5_tokenizer.tokenize(prompt)

src/mflux/flux/flux_initializer.py

@@ -0,0 +1,111 @@
+from mflux.controlnet.transformer_controlnet import TransformerControlnet
+from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
+from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
+from mflux.models.transformer.transformer import Transformer
+from mflux.models.vae.vae import VAE
+from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
+from mflux.tokenizer.t5_tokenizer import TokenizerT5
+from mflux.tokenizer.tokenizer_handler import TokenizerHandler
+from mflux.weights.weight_handler import WeightHandler
+from mflux.weights.weight_handler_lora import WeightHandlerLoRA
+from mflux.weights.weight_util import WeightUtil
+
+
+class FluxInitializer:
+    @staticmethod
+    def init(
+        flux_model,
+        model_config,
+        quantize: int | None,
+        local_path: str | None,
+        lora_paths: list[str] | None,
+        lora_scales: list[float] | None,
+    ) -> None:
+        # 0. Set paths and config for later
+        flux_model.lora_paths = lora_paths
+        flux_model.lora_scales = lora_scales
+        flux_model.model_config = model_config
+
+        # 1. Initialize tokenizers
+        tokenizers = TokenizerHandler(
+            repo_id=model_config.model_name,
+            max_t5_length=model_config.max_sequence_length,
+            local_path=local_path,
+        )
+        flux_model.t5_tokenizer = TokenizerT5(
+            tokenizer=tokenizers.t5,
+            max_length=model_config.max_sequence_length
+        )  # fmt: off
+        flux_model.clip_tokenizer = TokenizerCLIP(
+            tokenizer=tokenizers.clip,
+        )
+
+        # 2. Load the regular weights
+        weights = WeightHandler.load_regular_weights(
+            repo_id=model_config.model_name,
+            local_path=local_path
+        )  # fmt: off
+
+        # 3. Initialize all models
+        flux_model.vae = VAE()
+        flux_model.transformer = Transformer(
+            model_config=model_config,
+            num_transformer_blocks=weights.num_transformer_blocks(),
+            num_single_transformer_blocks=weights.num_single_transformer_blocks(),
+        )
+        flux_model.t5_text_encoder = T5Encoder()
+        flux_model.clip_text_encoder = CLIPEncoder()
+
+        # 4. Apply weights and quantize the models
+        flux_model.bits = WeightUtil.set_weights_and_quantize(
+            quantize_arg=quantize,
+            weights=weights,
+            vae=flux_model.vae,
+            transformer=flux_model.transformer,
+            t5_text_encoder=flux_model.t5_text_encoder,
+            clip_text_encoder=flux_model.clip_text_encoder,
+        )
+
+        # 5. Set LoRA weights
+        lora_weights = WeightHandlerLoRA.load_lora_weights(
+            transformer=flux_model.transformer,
+            lora_files=lora_paths,
+            lora_scales=lora_scales,
+        )
+        WeightHandlerLoRA.set_lora_weights(
+            transformer=flux_model.transformer,
+            loras=lora_weights
+        )  # fmt: off
+
+    @staticmethod
+    def init_controlnet(
+        flux_model,
+        model_config,
+        quantize: int | None,
+        local_path: str | None,
+        lora_paths: list[str] | None,
+        lora_scales: list[float] | None,
+    ) -> None:
+        # 1. Start with same init as regular Flux
+        FluxInitializer.init(
+            flux_model=flux_model,
+            model_config=model_config,
+            quantize=quantize,
+            local_path=local_path,
+            lora_paths=lora_paths,
+            lora_scales=lora_scales,
+        )
+
+        # 2. Apply ControlNet-specific initialization
+        weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer()
+        flux_model.transformer_controlnet = TransformerControlnet(
+            model_config=model_config,
+            num_transformer_blocks=weights_controlnet.num_transformer_blocks(),
+            num_single_transformer_blocks=weights_controlnet.num_single_transformer_blocks(),
+        )
+        WeightUtil.set_controlnet_weights_and_quantize(
+            quantize_arg=quantize,
+            weights=weights_controlnet,
+            transformer_controlnet=flux_model.transformer_controlnet,
+        )