diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py index 123bb44..1cc86b4 100644 --- a/src/mflux/config/config.py +++ b/src/mflux/config/config.py @@ -33,4 +33,4 @@ class ConfigControlnet(Config): controlnet_strength: float = 1.0, ): super().__init__(num_inference_steps, width, height, guidance) - self.controlnet_conditioning_scale = controlnet_strength + self.controlnet_strength = controlnet_strength diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index 67a357a..20205bc 100644 --- a/src/mflux/config/runtime_config.py +++ b/src/mflux/config/runtime_config.py @@ -37,7 +37,7 @@ class RuntimeConfig: return self.model_config.num_train_steps @property - def controlnet_conditioning_scale(self) -> float: + def controlnet_strength(self) -> float: if isinstance(self.config, ConfigControlnet): return self.config.controlnet_strength else: diff --git a/src/mflux/controlnet/__init__.py b/src/mflux/controlnet/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index a1aafac..2ed4c70 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -268,7 +268,7 @@ class TransformerControlnet(nn.Module): time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision) hidden_states = self.x_embedder(hidden_states) hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_cond) - conditioning_scale = config.config.controlnet_conditioning_scale + conditioning_scale = config.config.controlnet_strength guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)).astype(config.precision) text_embeddings = self.time_text_embed.forward(time_step, pooled_prompt_embeds, guidance) diff --git a/src/mflux/post_processing/image.py b/src/mflux/post_processing/image.py index fe69cd6..91aab40 100644 --- a/src/mflux/post_processing/image.py +++ b/src/mflux/post_processing/image.py @@ -26,6 +26,7 @@ class GeneratedImage: generation_time: float, lora_paths: list[str], lora_scales: list[float], + controlnet_strength: float | None = None, ): self.image = image self.model_config = model_config @@ -38,6 +39,7 @@ class GeneratedImage: self.generation_time = generation_time self.lora_paths = lora_paths self.lora_scales = lora_scales + self.controlnet_strength = controlnet_strength def save(self, path: str, export_json_metadata: bool = False) -> None: file_path = Path(path) @@ -123,4 +125,5 @@ class GeneratedImage: 'lora_paths': ', '.join(self.lora_paths) if self.lora_paths else '', 'lora_scales': ', '.join([f"{scale:.2f}" for scale in self.lora_scales]) if self.lora_scales else '', 'prompt': self.prompt, + 'controlnet_strength': "None" if self.controlnet_strength is None else f"{self.controlnet_strength:.2f}", } diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 2ef5b5e..dc54e7a 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -3,6 +3,7 @@ from PIL import Image import mlx.core as mx import numpy as np +from mflux.config.config import ConfigControlnet from mflux.config.runtime_config import RuntimeConfig from mflux.post_processing.image import GeneratedImage @@ -35,6 +36,7 @@ class ImageUtil: generation_time=generation_time, lora_paths=lora_paths, lora_scales=lora_scales, + controlnet_strength=config.controlnet_strength if isinstance(config.config, ConfigControlnet) else None, ) @staticmethod